Repository: chainer/chainer

Issue:
Same api for GRU and LSTM 
Currently, GRU is stateless while LSTM is stateful. It would be great to make GRU stateful as well, for the stake of consistency.



Discussion:
Thank you for suggestion. Just to be sure, do you agree with the idea that `links.LSTM` is stateful in the sense that it holds state (`self.h` and `self.c`) as attribute whereas `functions.LSTM` and `links.GRU` are stateless because they take current state (`h` and `c` for LSTM and `h` for GPU) as inputs and do not hold them as attributes?

One of the suggestion is, like LSTM, to make `links.GPU` stateful and `functions.GPU` stateless.

Yes, this is what I mean.

On Tue, Dec 22, 2015, 07:50 Kenta OONO notifications@github.com wrote:

> Thank you for suggestion. Just to be sure, do you agree with the idea that
> links.LSTM is stateful in the sense that it holds state (self.h and self.c)
> as attribute whereas functions.LSTM and links.GRU are stateless because
> they take current state (h and c for LSTM and h for GPU) as inputs and do
> not hold them as attributes?
> 
> One of the suggestion is, like LSTM, to make links.GPU stateful and
> functions.GPU stateless.
> 
> —
> Reply to this email directly or view it on GitHub
> https://github.com/pfnet/chainer/issues/789#issuecomment-166531761.

I see. I agree with your suggestion.


Resolving patch:
diff --git a/chainer/links/__init__.py b/chainer/links/__init__.py
--- a/chainer/links/__init__.py
+++ b/chainer/links/__init__.py
@@ -25,6 +25,7 @@
 Deconvolution2D = deconvolution_2d.Deconvolution2D
 EmbedID = embed_id.EmbedID
 GRU = gru.GRU
+StatefulGRU = gru.StatefulGRU
 Inception = inception.Inception
 InceptionBN = inceptionbn.InceptionBN
 Linear = linear.Linear
diff --git a/chainer/links/connection/gru.py b/chainer/links/connection/gru.py
--- a/chainer/links/connection/gru.py
+++ b/chainer/links/connection/gru.py
@@ -1,12 +1,30 @@
+import numpy
+
+import chainer
 from chainer.functions.activation import sigmoid
 from chainer.functions.activation import tanh
 from chainer import link
 from chainer.links.connection import linear
 
 
-class GRU(link.Chain):
+class GRUBase(link.Chain):
+
+    def __init__(self, n_units, n_inputs=None):
+        if n_inputs is None:
+            n_inputs = n_units
+        super(GRUBase, self).__init__(
+            W_r=linear.Linear(n_inputs, n_units),
+            U_r=linear.Linear(n_units, n_units),
+            W_z=linear.Linear(n_inputs, n_units),
+            U_z=linear.Linear(n_units, n_units),
+            W=linear.Linear(n_inputs, n_units),
+            U=linear.Linear(n_units, n_units),
+        )
+
+
+class GRU(GRUBase):
 
-    """Gated Recurrent Unit function (GRU).
+    """Stateless Gated Recurrent Unit function (GRU).
 
     GRU function has six parameters :math:`W_r`, :math:`W_z`, :math:`W`,
     :math:`U_r`, :math:`U_z`, and :math:`U`. All these parameters are
@@ -26,9 +44,14 @@ class GRU(link.Chain):
     where :math:`\\sigma` is the sigmoid function, and :math:`\\odot` is the
     element-wise product.
 
+    :class:`~chainer.links.GRU` does not hold the value of
+    hidden vector :math:`h`. So this is *stateless*.
+    Use :class:`~chainer.links.StatefulGRU` as a *stateful* GRU.
+
     Args:
-        n_units(int): Dimension of input vector :math:`x`, and hidden vector
-            :math:`h`.
+        n_units(int): Dimension of hidden vector :math:`h`.
+        n_inputs(int): Dimension of input vector :math:`x`. If ``None``,
+        it is set to the same value as ``n_units``.
 
     See:
         - `On the Properties of Neural Machine Translation: Encoder-Decoder
@@ -38,17 +61,9 @@ class GRU(link.Chain):
           Modeling <http://arxiv.org/abs/1412.3555>`_
           [Chung+NIPS2014 DLWorkshop].
 
-    """
 
-    def __init__(self, n_units):
-        super(GRU, self).__init__(
-            W_r=linear.Linear(n_units, n_units),
-            U_r=linear.Linear(n_units, n_units),
-            W_z=linear.Linear(n_units, n_units),
-            U_z=linear.Linear(n_units, n_units),
-            W=linear.Linear(n_units, n_units),
-            U=linear.Linear(n_units, n_units),
-        )
+    .. seealso:: :class:`~chainer.links.StatefulGRU`
+    """
 
     def __call__(self, h, x):
         r = sigmoid.sigmoid(self.W_r(x) + self.U_r(h))
@@ -56,3 +71,82 @@ def __call__(self, h, x):
         h_bar = tanh.tanh(self.W(x) + self.U(r * h))
         h_new = (1 - z) * h + z * h_bar
         return h_new
+
+
+class StatefulGRU(GRUBase):
+    """Stateful Gated Recurrent Unit function (GRU).
+
+    Stateful GRU function has six parameters :math:`W_r`, :math:`W_z`,
+    :math:`W`, :math:`U_r`, :math:`U_z`, and :math:`U`.
+    All these parameters are :math:`n \\times n` matricies,
+    where :math:`n` is the dimension of hidden vectors.
+
+    Given input vector :math:`x`, Stateful GRU returns the next
+    hidden vector :math:`h'` defined as
+
+    .. math::
+
+       r &=& \\sigma(W_r x + U_r h), \\\\
+       z &=& \\sigma(W_z x + U_z h), \\\\
+       \\bar{h} &=& \\tanh(W x + U (r \\odot h)), \\\\
+       h' &=& (1 - z) \\odot h + z \\odot \\bar{h},
+
+    where :math:`h` is current hidden vector.
+
+    As the name indicates, :class:`~chainer.links.StatefulGRU` is *stateful*,
+    meaning that it also holds the next hidden vector `h'` as a state.
+    Use :class:`~chainer.links.GRU` as a stateless version of GRU.
+
+    Args:
+        in_size(int): Dimension of input vector :math:`x`.
+        out_size(int): Dimension of hidden vector :math:`h`.
+
+    Attributes:
+        h(~chainer.Variable): Hidden vector that indicates the state of
+        :class:`~chainer.links.StatefulGRU`.
+
+    .. seealso:: :class:`~chainer.functions.GRU`
+    """
+
+    def __init__(self, in_size, out_size):
+        super(StatefulGRU, self).__init__(out_size, in_size)
+        self.state_size = out_size
+        self.reset_state()
+
+    def to_cpu(self):
+        super(StatefulGRU, self).to_cpu()
+        if self.h is not None:
+            self.h.to_cpu()
+
+    def to_gpu(self, device=None):
+        super(StatefulGRU, self).to_gpu(device)
+        if self.h is not None:
+            self.h.to_gpu(device)
+
+    def set_state(self, h):
+        assert isinstance(h, chainer.Variable)
+        h_ = h
+        if self.xp == numpy:
+            h_.to_cpu()
+        else:
+            h_.to_gpu()
+        self.h = h_
+
+    def reset_state(self):
+        self.h = None
+
+    def __call__(self, x):
+        z = self.W_z(x)
+        h_bar = self.W(x)
+        if self.h is not None:
+            r = sigmoid.sigmoid(self.W_r(x) + self.U_r(self.h))
+            z += self.U_z(self.h)
+            h_bar += self.U(r * self.h)
+        z = sigmoid.sigmoid(z)
+        h_bar = tanh.tanh(h_bar)
+
+        h_new = z * h_bar
+        if self.h is not None:
+            h_new += (1 - z) * self.h
+        self.h = h_new
+        return self.h


Test patch:
diff --git a/tests/chainer_tests/links_tests/connection_tests/test_gru.py b/tests/chainer_tests/links_tests/connection_tests/test_gru.py
--- a/tests/chainer_tests/links_tests/connection_tests/test_gru.py
+++ b/tests/chainer_tests/links_tests/connection_tests/test_gru.py
@@ -25,22 +25,57 @@ def _gru(func, h, x):
     return y
 
 
+@testing.parameterize(
+    {'gru': links.GRU, 'state': 'random', 'in_size': 4, 'out_size': 8},
+    {'gru': links.GRU, 'state': 'random', 'out_size': 8},
+    {'gru': links.StatefulGRU, 'state': 'random', 'in_size': 4, 'out_size': 8},
+    {'gru': links.StatefulGRU, 'state': 'zero', 'in_size': 4, 'out_size': 8},
+)
 class TestGRU(unittest.TestCase):
 
     def setUp(self):
-        self.link = links.GRU(8)
-        self.x = numpy.random.uniform(-1, 1, (3, 8)).astype(numpy.float32)
-        self.h = numpy.random.uniform(-1, 1, (3, 8)).astype(numpy.float32)
-        self.gy = numpy.random.uniform(-1, 1, (3, 8)).astype(numpy.float32)
+        if self.gru == links.GRU:
+            if hasattr(self, 'in_size'):
+                self.link = self.gru(self.out_size, self.in_size)
+            else:
+                self.link = self.gru(self.out_size)
+                self.in_size = self.out_size
+        elif self.gru == links.StatefulGRU:
+            self.link = self.gru(self.in_size, self.out_size)
+        else:
+            self.fail('Unsupported link(only GRU and StatefulGRU '
+                      'are supported):{}'.format(self.gru))
+
+        self.x = numpy.random.uniform(
+            -1, 1, (3, self.in_size)).astype(numpy.float32)
+        if self.state == 'random':
+            self.h = numpy.random.uniform(
+                -1, 1, (3, self.out_size)).astype(numpy.float32)
+        elif self.state == 'zero':
+            self.h = numpy.zeros((3, self.out_size), dtype=numpy.float32)
+        else:
+            self.fail('Unsupported state initialization:{}'.format(self.state))
+        self.gy = numpy.random.uniform(
+            -1, 1, (3, self.out_size)).astype(numpy.float32)
+
+    def _forward(self, link, h, x):
+        if isinstance(link, links.GRU):
+            return link(h, x)
+        else:
+            if self.state != 'zero':
+                link.set_state(h)
+            return link(x)
 
     def check_forward(self, h_data, x_data):
         h = chainer.Variable(h_data)
         x = chainer.Variable(x_data)
-        y = self.link(h, x)
-        self.assertEqual(y.data.dtype, numpy.float32)
+        y = self._forward(self.link, h, x)
 
+        self.assertEqual(y.data.dtype, numpy.float32)
         y_expect = _gru(self.link, h_data, x_data)
         gradient_check.assert_allclose(y_expect, y.data)
+        if isinstance(self.link, links.StatefulGRU):
+            gradient_check.assert_allclose(self.link.h.data, y.data)
 
     def test_forward_cpu(self):
         self.check_forward(self.h, self.x)
@@ -54,17 +89,18 @@ def test_forward_gpu(self):
     def check_backward(self, h_data, x_data, y_grad):
         h = chainer.Variable(h_data)
         x = chainer.Variable(x_data)
-        y = self.link(h, x)
-
+        y = self._forward(self.link, h, x)
         y.grad = y_grad
         y.backward()
 
-        f = lambda: (_gru(self.link, h.data, x.data),)
-        gh, gx = gradient_check.numerical_grad(f, (h.data, x.data), (y.grad,))
-
-        gradient_check.assert_allclose(gh, h.grad, atol=1e-3)
+        f = lambda: (_gru(self.link, h_data, x_data),)
+        gx, = gradient_check.numerical_grad(f, (x.data,), (y.grad,))
         gradient_check.assert_allclose(gx, x.grad, atol=1e-3)
 
+        if isinstance(self.link, links.GRU):
+            gh, = gradient_check.numerical_grad(f, (h.data,), (y.grad,))
+            gradient_check.assert_allclose(gh, h.grad, atol=1e-3)
+
     def test_backward_cpu(self):
         self.check_backward(self.h, self.x, self.gy)
 
@@ -76,4 +112,87 @@ def test_backward_gpu(self):
                             cuda.to_gpu(self.gy))
 
 
+@testing.parameterize(
+    *testing.product({
+        'link_array_module': ['to_cpu', 'to_gpu'],
+        'state_array_module': ['to_cpu', 'to_gpu']
+    }))
+class TestGRUState(unittest.TestCase):
+
+    def setUp(self):
+        in_size, out_size = 10, 8
+        self.link = links.StatefulGRU(in_size, out_size)
+        self.h = chainer.Variable(
+            numpy.random.uniform(-1, 1, (3, out_size)).astype(numpy.float32))
+
+    def check_set_state(self, h):
+        self.link.set_state(h)
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+
+    def test_set_state_cpu(self):
+        self.check_set_state(self.h)
+
+    @attr.gpu
+    def test_set_state_gpu(self):
+        getattr(self.link, self.link_array_module)()
+        getattr(self.h, self.state_array_module)()
+        self.check_set_state(self.h)
+
+    def check_reset_state(self):
+        self.link.reset_state()
+        self.assertIsNone(self.link.h)
+
+    def test_reset_state_cpu(self):
+        self.check_reset_state()
+
+    @attr.gpu
+    def test_reset_state_gpu(self):
+        getattr(self.link, self.link_array_module)()
+        self.check_reset_state()
+
+
+class TestGRUToCPUToGPU(unittest.TestCase):
+
+    def setUp(self):
+        in_size, out_size = 10, 8
+        self.link = links.StatefulGRU(in_size, out_size)
+        self.h = chainer.Variable(
+            numpy.random.uniform(-1, 1, (3, out_size)).astype(numpy.float32))
+
+    def check_to_cpu(self, h):
+        self.link.set_state(h)
+        self.link.to_cpu()
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+        self.link.to_cpu()
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+
+    def test_to_cpu_cpu(self):
+        self.check_to_cpu(self.h)
+
+    @attr.gpu
+    def test_to_cpu_gpu(self):
+        self.h.to_gpu()
+        self.check_to_cpu(self.h)
+
+    def check_to_cpu_to_gpu(self, h):
+        self.link.set_state(h)
+        self.link.to_gpu()
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+        self.link.to_gpu()
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+        self.link.to_cpu()
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+        self.link.to_gpu()
+        self.assertIsInstance(self.link.h.data, self.link.xp.ndarray)
+
+    @attr.gpu
+    def test_to_cpu_to_gpu_cpu(self):
+        self.check_to_cpu_to_gpu(self.h)
+
+    @attr.gpu
+    def test_to_cpu_to_gpu_gpu(self):
+        self.h.to_gpu()
+        self.check_to_cpu_to_gpu(self.h)
+
+
 testing.run_module(__name__, __file__)