Python API for StaticRNN (#4991)

Author: reyoung

python/paddle/v2/framework/framework.py

@@ -113,6 +113,10 @@ def data_type(self):
     def lod_level(self):
         return self.desc.lod_level()
 
+    @property
+    def type(self):
+        return self.desc.type()
+
     @staticmethod
     def _unique_var_name_():
         uid = core.unique_integer()  # unique during whole process.

python/paddle/v2/framework/layer_helper.py

@@ -1,8 +1,11 @@
-from paddle.v2.framework.framework import Variable, OpProtoHolder, g_program, g_init_program
-import paddle.v2.framework.core as core
 import copy
 import itertools
 
+import paddle.v2.framework.core as core
+
+from paddle.v2.framework.framework import Variable, g_program, \
+    g_init_program
+
 
 def unique_name(prefix):
     uid = core.unique_integer()  # unique during whole process.
@@ -130,6 +133,9 @@ def create_tmp_variable(self, dtype):
             dtype=dtype,
             persistable=False)
 
+    def create_variable(self, *args, **kwargs):
+        return self.program.current_block().create_var(*args, **kwargs)
+
     def create_global_variable(self, *args, **kwargs):
         return self.program.global_block().create_var(
             *args, persistable=False, **kwargs)

python/paddle/v2/framework/layers.py

@@ -1,10 +1,11 @@
-from paddle.v2.framework.layer_helper import LayerHelper
+from paddle.v2.framework.layer_helper import LayerHelper, unique_name
 import paddle.v2.framework.core as core
-from paddle.v2.framework.framework import OpProtoHolder, Variable
+from paddle.v2.framework.framework import OpProtoHolder, Variable, Program
 import re
 
 __all__ = [
-    'fc', 'data', 'cross_entropy', 'conv2d', 'pool2d', 'embedding', 'concat'
+    'fc', 'data', 'cross_entropy', 'conv2d', 'pool2d', 'embedding', 'concat',
+    'StaticRNN'
 ]
 
 
@@ -26,7 +27,9 @@ def fc(input,
     mul_results = []
     for input_var, param_attr in helper.iter_inputs_and_params():
         input_shape = input_var.shape
-        param_shape = list(input_shape[num_flatten_dims:]) + [size]
+        param_shape = [
+            reduce(lambda a, b: a * b, input_shape[num_flatten_dims:], 1)
+        ] + [size]
 
         w = helper.create_parameter(
             attr=param_attr, shape=param_shape, dtype=dtype)
@@ -38,10 +41,8 @@ def fc(input,
                 "Y": w,
             },
             outputs={"Out": tmp},
-            attrs={
-                'x_num_col_dims': num_flatten_dims,
-                'y_num_col_dims': len(input_shape) - num_flatten_dims
-            })
+            attrs={'x_num_col_dims': num_flatten_dims,
+                   'y_num_col_dims': 1})
         mul_results.append(tmp)
 
     # sum
@@ -273,3 +274,170 @@ def pool2d(input,
         })
 
     return pool_out
+
+
+class BlockGuard(object):
+    """
+    BlockGuard used to create sub-block in program by using Python `with` 
+    keyword.
+    """
+
+    def __init__(self, program):
+        if not isinstance(program, Program):
+            raise TypeError("BlockGuard takes a program")
+        self.program = program
+
+    def __enter__(self):
+        self.program.create_block()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.program.rollback()
+        if exc_type is not None:
+            return False  # re-raise exception
+        return True
+
+
+class StaticRNNGuard(BlockGuard):
+    def __init__(self, rnn):
+        if not isinstance(rnn, StaticRNN):
+            raise TypeError("StaticRNNGuard takes an StaticRNN")
+        super(StaticRNNGuard, self).__init__(rnn.helper.program)
+        self.rnn = rnn
+
+    def __enter__(self):
+        self.rnn.status = StaticRNN.IN_RNN_BLOCK
+        return super(StaticRNNGuard, self).__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.rnn.status = StaticRNN.AFTER_RNN_BLOCK
+        self.rnn.complete_rnn_op()
+        return super(StaticRNNGuard, self).__exit__(exc_type, exc_val, exc_tb)
+
+
+class StaticRNNMemoryLink(object):
+    """
+    :param init: the initial variable for Memory
+    :type init: Variable
+    :param pre_mem: the memory variable in previous time step
+    :type pre_mem: Variable
+    :param mem: the memory variable in current time step
+    :type mem: Variable
+    """
+
+    def __init__(self, init, pre_mem, mem=None):
+        self.init = init
+        self.pre_mem = pre_mem
+        self.mem = mem
+
+
+class StaticRNN(object):
+    BEFORE_RNN_BLOCK = 0
+    IN_RNN_BLOCK = 1
+    AFTER_RNN_BLOCK = 2
+
+    def __init__(self, name=None, program=None):
+        self.helper = LayerHelper("static_rnn", name=name, program=program)
+        self.memories = {}  # memory map, from pre_mem.name --> MemoryLink
+        self.inputs = []  # input variable list in current block
+        self.outputs = []  # output variable list in parent block
+        self.status = StaticRNN.BEFORE_RNN_BLOCK  # status flag.
+        # sequence length, since it is a static RNN, sequence length are fixed.
+        self.seq_len = None
+
+    def step(self):
+        return StaticRNNGuard(self)
+
+    def _assert_in_rnn_block_(self, method):
+        if self.status != StaticRNN.IN_RNN_BLOCK:
+            raise ValueError("You must invoke {0} in rnn block".format(method))
+
+    def memory(self, init=None, shape=None, dtype=None, init_value=0):
+        self._assert_in_rnn_block_('memory')
+        if init is None:
+            if shape is None or dtype is None:
+                raise ValueError(
+                    "if init is None, memory at least need shape and dtype")
+            parent_block = self.parent_block()
+            var_name = unique_name("@".join([self.helper.name, "memory_boot"]))
+            boot_var = parent_block.create_var(
+                name=var_name, shape=shape, dtype=dtype, persistable=False)
+
+            parent_block.append_op(
+                type="fill_constant",
+                inputs={},
+                outputs={'Out': [boot_var]},
+                attrs={
+                    'value': init_value,
+                    'shape': boot_var.shape,
+                    'data_type': boot_var.data_type
+                })
+
+            return self.memory(init=boot_var)
+        else:
+            pre_mem = self.helper.create_variable(
+                name=unique_name("@".join([self.helper.name, "mem"])),
+                dtype=init.data_type,
+                shape=init.shape)
+            self.memories[pre_mem.name] = StaticRNNMemoryLink(
+                init=init, pre_mem=pre_mem)
+            return pre_mem
+
+    def step_input(self, x):
+        self._assert_in_rnn_block_('step_input')
+        if not isinstance(x, Variable):
+            raise TypeError("step input takes a Variable")
+        if self.seq_len is None:
+            self.seq_len = x.shape[1]
+        elif self.seq_len != x.shape[1]:
+            raise ValueError("Static RNN only take fix seq_len input")
+
+        ipt = self.helper.create_variable(
+            name=x.name,
+            dtype=x.data_type,
+            shape=[-1] + list(x.shape[2:]),
+            type=x.type)
+        self.inputs.append(ipt)
+        return ipt
+
+    def step_output(self, o):
+        self._assert_in_rnn_block_('step_output')
+        if not isinstance(o, Variable):
+            raise TypeError("step output takes a Variable")
+
+        out_var = self.parent_block().create_var(
+            name=o.name,
+            shape=[-1, self.seq_len] + list(o.shape[1:]),
+            dtype=o.data_type)
+
+        self.outputs.append(out_var)
+
+    def output(self, *outputs):
+        for each in outputs:
+            self.step_output(each)
+
+    def update_memory(self, mem, var):
+        if not isinstance(mem, Variable) or not isinstance(var, Variable):
+            raise TypeError("update memory should take variables")
+        self.memories[mem.name].mem = var
+
+    def parent_block(self):
+        prog = self.helper.program
+        parent_idx = prog.current_block().parent_idx
+        assert parent_idx >= 0
+        parent_block = prog.block(parent_idx)
+        return parent_block
+
+    def __call__(self, *args, **kwargs):
+        if self.status != StaticRNN.AFTER_RNN_BLOCK:
+            raise ValueError("RNN output can only be retrieved after rnn block")
+        if len(self.outputs) == 0:
+            raise ValueError("RNN has no output")
+        elif len(self.outputs) == 1:
+            return self.outputs[0]
+        else:
+            return self.outputs
+
+    def complete_rnn_op(self):
+        # TODO(yuyang18): Create RNN Op here.
+        # Implement this method after RNN op complete.
+        pass

python/paddle/v2/framework/tests/test_rnn_helpers.py

@@ -0,0 +1,38 @@
+import unittest
+from paddle.v2.framework.layers import *
+from paddle.v2.framework.framework import g_program
+
+
+class TestRNN(unittest.TestCase):
+    def test_rnn(self):
+        img = data(
+            shape=[
+                80,  # sequence length
+                22,  # image height
+                22
+            ],  # image width
+            data_type='float32',
+            name='image')
+        hidden = fc(input=img, size=100, act='sigmoid', num_flatten_dims=2)
+        self.assertEqual((-1, 80, 100), hidden.shape)
+        hidden = fc(input=hidden, size=100, act='sigmoid', num_flatten_dims=2)
+        self.assertEqual((-1, 80, 100), hidden.shape)
+
+        rnn = StaticRNN()
+        with rnn.step():
+            hidden = rnn.step_input(hidden)
+            self.assertEqual((-1, 100), hidden.shape)
+            memory = rnn.memory(shape=(-1, 32), dtype='float32', init_value=0.0)
+
+            rnn_out = fc(input=[hidden, memory], size=32, act='sigmoid')
+            self.assertEqual((-1, 32), rnn_out.shape)
+            rnn.update_memory(memory, rnn_out)
+            rnn.output(rnn_out)
+
+        out = rnn()
+        self.assertEqual((-1, 80, 32), out.shape)
+        print g_program
+
+
+if __name__ == '__main__':
+    unittest.main()