Repository: chainer/chainer

Issue:
Refactor NStepRNNs' links
All NStepRNNs have similar implementation. It is difficult to add new feature for them. We can merge it.


Discussion:

Resolving patch:
diff --git a/chainer/links/connection/n_step_gru.py b/chainer/links/connection/n_step_gru.py
--- a/chainer/links/connection/n_step_gru.py
+++ b/chainer/links/connection/n_step_gru.py
@@ -1,19 +1,8 @@
-import numpy
-import six
-
-from chainer.backends import cuda
-from chainer.functions.array import permutate
-from chainer.functions.array import transpose_sequence
 from chainer.functions.connection import n_step_gru as rnn
-from chainer.initializers import normal
-from chainer import link
-from chainer.links.connection.n_step_rnn import argsort_list_descent
-from chainer.links.connection.n_step_rnn import permutate_list
-from chainer.utils import argument
-from chainer import variable
+from chainer.links.connection import n_step_rnn
 
 
-class NStepGRUBase(link.ChainList):
+class NStepGRUBase(n_step_rnn.NStepRNNBase):
 
     """__init__(self, n_layers, in_size, out_size, dropout, use_bi_direction)
 
@@ -42,96 +31,7 @@ class NStepGRUBase(link.ChainList):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, use_bi_direction,
-                 **kwargs):
-        argument.check_unexpected_kwargs(
-            kwargs, use_cudnn='use_cudnn argument is not supported anymore. '
-            'Use chainer.using_config')
-        argument.assert_kwargs_empty(kwargs)
-
-        weights = []
-        direction = 2 if use_bi_direction else 1
-        for i in six.moves.range(n_layers):
-            for di in six.moves.range(direction):
-                weight = link.Link()
-                with weight.init_scope():
-                    for j in six.moves.range(6):
-                        if i == 0 and j < 3:
-                            w_in = in_size
-                        elif i > 0 and j < 3:
-                            w_in = out_size * direction
-                        else:
-                            w_in = out_size
-                        w = variable.Parameter(
-                            normal.Normal(numpy.sqrt(1. / w_in)),
-                            (out_size, w_in))
-                        b = variable.Parameter(0, (out_size,))
-                        setattr(weight, 'w%d' % j, w)
-                        setattr(weight, 'b%d' % j, b)
-                weights.append(weight)
-
-        super(NStepGRUBase, self).__init__(*weights)
-
-        self.n_layers = n_layers
-        self.dropout = dropout
-        self.out_size = out_size
-        self.direction = direction
-        self.rnn = rnn.n_step_bigru if use_bi_direction else rnn.n_step_gru
-
-    def init_hx(self, xs):
-        shape = (self.n_layers * self.direction, len(xs), self.out_size)
-        with cuda.get_device_from_id(self._device_id):
-            hx = variable.Variable(self.xp.zeros(shape, dtype=xs[0].dtype))
-        return hx
-
-    def __call__(self, hx, xs, **kwargs):
-        """__call__(self, hx, xs)
-
-        Calculate all hidden states and cell states.
-
-        .. warning::
-
-           ``train`` argument is not supported anymore since v2.
-           Instead, use ``chainer.using_config('train', train)``.
-           See :func:`chainer.using_config`.
-
-        Args:
-            hx (~chainer.Variable or None): Initial hidden states. If ``None``
-                is specified zero-vector is used.
-            xs (list of ~chainer.Variable): List of input sequences.
-                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
-                a sequence.
-
-        """
-        argument.check_unexpected_kwargs(
-            kwargs, train='train argument is not supported anymore. '
-            'Use chainer.using_config')
-        argument.assert_kwargs_empty(kwargs)
-
-        assert isinstance(xs, (list, tuple))
-        xp = cuda.get_array_module(hx, *xs)
-        indices = argsort_list_descent(xs)
-        indices_array = xp.array(indices)
-
-        xs = permutate_list(xs, indices, inv=False)
-        if hx is None:
-            hx = self.init_hx(xs)
-        else:
-            hx = permutate.permutate(hx, indices_array, axis=1, inv=False)
-
-        trans_x = transpose_sequence.transpose_sequence(xs)
-
-        ws = [[w.w0, w.w1, w.w2, w.w3, w.w4, w.w5] for w in self]
-        bs = [[w.b0, w.b1, w.b2, w.b3, w.b4, w.b5] for w in self]
-
-        hy, trans_y = self.rnn(
-            self.n_layers, self.dropout, hx, ws, bs, trans_x)
-
-        hy = permutate.permutate(hy, indices_array, axis=1, inv=True)
-        ys = transpose_sequence.transpose_sequence(trans_y)
-        ys = permutate_list(ys, indices, inv=True)
-
-        return hy, ys
+    n_weights = 6
 
 
 class NStepGRU(NStepGRUBase):
@@ -166,10 +66,10 @@ class NStepGRU(NStepGRUBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
-        NStepGRUBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            use_bi_direction=False, **kwargs)
+    use_bi_direction = False
+
+    def rnn(self, *args):
+        return rnn.n_step_gru(*args)
 
 
 class NStepBiGRU(NStepGRUBase):
@@ -204,7 +104,7 @@ class NStepBiGRU(NStepGRUBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
-        NStepGRUBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            use_bi_direction=True, **kwargs)
+    use_bi_direction = True
+
+    def rnn(self, *args):
+        return rnn.n_step_bigru(*args)
diff --git a/chainer/links/connection/n_step_lstm.py b/chainer/links/connection/n_step_lstm.py
--- a/chainer/links/connection/n_step_lstm.py
+++ b/chainer/links/connection/n_step_lstm.py
@@ -1,17 +1,8 @@
-import six
-
-from chainer.backends import cuda
-from chainer.functions.array import permutate
-from chainer.functions.array import transpose_sequence
 from chainer.functions.connection import n_step_lstm as rnn
-from chainer import initializers
-from chainer import link
 from chainer.links.connection import n_step_rnn
-from chainer.utils import argument
-from chainer import variable
 
 
-class NStepLSTMBase(link.ChainList):
+class NStepLSTMBase(n_step_rnn.NStepRNNBase):
     """Base link class for Stacked LSTM/BiLSTM links.
 
     This link is base link class for :func:`chainer.links.NStepLSTM` and
@@ -38,52 +29,7 @@ class NStepLSTMBase(link.ChainList):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout,
-                 initialW, initial_bias, use_bi_direction,
-                 **kwargs):
-        argument.check_unexpected_kwargs(
-            kwargs, use_cudnn='use_cudnn argument is not supported anymore. '
-            'Use chainer.using_config')
-        argument.assert_kwargs_empty(kwargs)
-
-        if initial_bias is None:
-            initial_bias = initializers.constant.Zero()
-        initialW = initializers._get_initializer(initialW)
-
-        weights = []
-        direction = 2 if use_bi_direction else 1
-        for i in six.moves.range(n_layers):
-            for di in six.moves.range(direction):
-                weight = link.Link()
-                with weight.init_scope():
-                    for j in six.moves.range(8):
-                        if i == 0 and j < 4:
-                            w_in = in_size
-                        elif i > 0 and j < 4:
-                            w_in = out_size * direction
-                        else:
-                            w_in = out_size
-                        name_w = 'w{}'.format(j)
-                        name_b = 'b{}'.format(j)
-                        w = variable.Parameter(initialW, (out_size, w_in))
-                        b = variable.Parameter(initial_bias, (out_size,))
-                        setattr(weight, name_w, w)
-                        setattr(weight, name_b, b)
-                weights.append(weight)
-
-        super(NStepLSTMBase, self).__init__(*weights)
-
-        self.n_layers = n_layers
-        self.dropout = dropout
-        self.out_size = out_size
-        self.direction = direction
-        self.rnn = rnn.n_step_bilstm if use_bi_direction else rnn.n_step_lstm
-
-    def init_hx(self, xs):
-        shape = (self.n_layers * self.direction, len(xs), self.out_size)
-        with cuda.get_device_from_id(self._device_id):
-            hx = variable.Variable(self.xp.zeros(shape, dtype=xs[0].dtype))
-        return hx
+    n_weights = 8
 
     def __call__(self, hx, cx, xs, **kwargs):
         """__call__(self, hx, cx, xs)
@@ -105,40 +51,7 @@ def __call__(self, hx, cx, xs, **kwargs):
                 Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                 a sequence.
         """
-        argument.check_unexpected_kwargs(
-            kwargs, train='train argument is not supported anymore. '
-            'Use chainer.using_config')
-        argument.assert_kwargs_empty(kwargs)
-
-        assert isinstance(xs, (list, tuple))
-        xp = cuda.get_array_module(hx, *xs)
-        indices = n_step_rnn.argsort_list_descent(xs)
-        indices_array = xp.array(indices)
-
-        xs = n_step_rnn.permutate_list(xs, indices, inv=False)
-        if hx is None:
-            hx = self.init_hx(xs)
-        else:
-            hx = permutate.permutate(hx, indices_array, axis=1, inv=False)
-
-        if cx is None:
-            cx = self.init_hx(xs)
-        else:
-            cx = permutate.permutate(cx, indices_array, axis=1, inv=False)
-
-        trans_x = transpose_sequence.transpose_sequence(xs)
-
-        ws = [[w.w0, w.w1, w.w2, w.w3, w.w4, w.w5, w.w6, w.w7] for w in self]
-        bs = [[w.b0, w.b1, w.b2, w.b3, w.b4, w.b5, w.b6, w.b7] for w in self]
-
-        hy, cy, trans_y = self.rnn(
-            self.n_layers, self.dropout, hx, cx, ws, bs, trans_x)
-
-        hy = permutate.permutate(hy, indices_array, axis=1, inv=True)
-        cy = permutate.permutate(cy, indices_array, axis=1, inv=True)
-        ys = transpose_sequence.transpose_sequence(trans_y)
-        ys = n_step_rnn.permutate_list(ys, indices, inv=True)
-
+        (hy, cy), ys = self._call([hx, cx], xs, **kwargs)
         return hy, cy, ys
 
 
@@ -179,12 +92,10 @@ class NStepLSTM(NStepLSTMBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout,
-                 initialW=None, initial_bias=None, **kwargs):
-        NStepLSTMBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            initialW, initial_bias,
-            use_bi_direction=False, **kwargs)
+    use_bi_direction = False
+
+    def rnn(self, *args):
+        return rnn.n_step_lstm(*args)
 
 
 class NStepBiLSTM(NStepLSTMBase):
@@ -224,9 +135,7 @@ class NStepBiLSTM(NStepLSTMBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout,
-                 initialW=None, initial_bias=None, **kwargs):
-        NStepLSTMBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            initialW, initial_bias,
-            use_bi_direction=True, **kwargs)
+    use_bi_direction = True
+
+    def rnn(self, *args):
+        return rnn.n_step_bilstm(*args)
diff --git a/chainer/links/connection/n_step_rnn.py b/chainer/links/connection/n_step_rnn.py
--- a/chainer/links/connection/n_step_rnn.py
+++ b/chainer/links/connection/n_step_rnn.py
@@ -27,7 +27,7 @@ def permutate_list(lst, indices, inv):
 
 
 class NStepRNNBase(link.ChainList):
-    """__init__(self, n_layers, in_size, out_size, dropout, use_bi_direction, activation)
+    """__init__(self, n_layers, in_size, out_size, dropout)
 
     Base link class for Stacked RNN/BiRNN links.
 
@@ -47,10 +47,6 @@ class NStepRNNBase(link.ChainList):
         in_size (int): Dimensionality of input vectors.
         out_size (int): Dimensionality of hidden states and output vectors.
         dropout (float): Dropout ratio.
-        use_bi_direction (bool): if ``True``, use Bi-directional RNN.
-            if ``False``, use Uni-directional RNN.
-        activation (str): Activation function name.
-            Please select ``tanh`` or ``relu``.
 
     .. seealso::
         :func:`chainer.links.NStepRNNReLU`
@@ -60,23 +56,28 @@ class NStepRNNBase(link.ChainList):
 
     """  # NOQA
 
-    def __init__(self, n_layers, in_size, out_size, dropout,
-                 use_bi_direction, activation, **kwargs):
+    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
         argument.check_unexpected_kwargs(
             kwargs, use_cudnn='use_cudnn argument is not supported anymore. '
-            'Use chainer.using_config')
+            'Use chainer.using_config',
+            use_bi_direction='use_bi_direction is not supported anymore',
+            activation='activation is not supported anymore')
         argument.assert_kwargs_empty(kwargs)
 
         weights = []
-        direction = 2 if use_bi_direction else 1
+        if self.use_bi_direction:
+            direction = 2
+        else:
+            direction = 1
+
         for i in six.moves.range(n_layers):
             for di in six.moves.range(direction):
                 weight = link.Link()
                 with weight.init_scope():
-                    for j in six.moves.range(2):
-                        if i == 0 and j < 1:
+                    for j in six.moves.range(self.n_weights):
+                        if i == 0 and j < self.n_weights // 2:
                             w_in = in_size
-                        elif i > 0 and j < 1:
+                        elif i > 0 and j < self.n_weights // 2:
                             w_in = out_size * direction
                         else:
                             w_in = out_size
@@ -90,12 +91,17 @@ def __init__(self, n_layers, in_size, out_size, dropout,
 
         super(NStepRNNBase, self).__init__(*weights)
 
+        self.ws = [[getattr(layer, 'w%d' % i)
+                    for i in six.moves.range(self.n_weights)]
+                   for layer in self]
+        self.bs = [[getattr(layer, 'b%d' % i)
+                    for i in six.moves.range(self.n_weights)]
+                   for layer in self]
+
         self.n_layers = n_layers
         self.dropout = dropout
-        self.activation = activation
         self.out_size = out_size
         self.direction = direction
-        self.rnn = rnn.n_step_birnn if use_bi_direction else rnn.n_step_rnn
 
     def init_hx(self, xs):
         shape = (self.n_layers * self.direction, len(xs), self.out_size)
@@ -103,6 +109,13 @@ def init_hx(self, xs):
             hx = variable.Variable(self.xp.zeros(shape, dtype=xs[0].dtype))
         return hx
 
+    def rnn(self, *args):
+        """Calls RNN function.
+
+        This function must be implemented in a child class.
+        """
+        raise NotImplementedError
+
     def __call__(self, hx, xs, **kwargs):
         """__call__(self, hx, xs)
 
@@ -121,36 +134,55 @@ def __call__(self, hx, xs, **kwargs):
                 Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                 a sequence.
         """
+        (hy,), ys = self._call([hx], xs, **kwargs)
+        return hy, ys
+
+    def _call(self, hs, xs, **kwargs):
+        """Calls RNN function.
+
+        Args:
+            hs (list of ~chainer.Variable or None): Lisit of hidden states.
+                Its length depends on its implementation.
+                If ``None`` is specified zero-vector is used.
+            xs (list of ~chainer.Variable): List of input sequences.
+                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
+                a sequence.
+
+        Returns:
+            tuple: hs
+        """
         argument.check_unexpected_kwargs(
             kwargs, train='train argument is not supported anymore. '
             'Use chainer.using_config')
         argument.assert_kwargs_empty(kwargs)
 
         assert isinstance(xs, (list, tuple))
-        xp = cuda.get_array_module(hx, *xs)
+        xp = cuda.get_array_module(*(list(hs) + list(xs)))
         indices = argsort_list_descent(xs)
         indices_array = xp.array(indices)
 
         xs = permutate_list(xs, indices, inv=False)
-        if hx is None:
-            hx = self.init_hx(xs)
-        else:
-            hx = permutate.permutate(hx, indices_array, axis=1, inv=False)
+        hxs = []
+        for hx in hs:
+            if hx is None:
+                hx = self.init_hx(xs)
+            else:
+                hx = permutate.permutate(hx, indices_array, axis=1, inv=False)
+            hxs.append(hx)
 
         trans_x = transpose_sequence.transpose_sequence(xs)
 
-        ws = [[w.w0, w.w1] for w in self]
-        bs = [[w.b0, w.b1] for w in self]
+        args = [self.n_layers, self.dropout] + hxs + \
+               [self.ws, self.bs, trans_x]
+        result = self.rnn(*args)
 
-        hy, trans_y = self.rnn(
-            self.n_layers, self.dropout, hx, ws, bs, trans_x,
-            activation=self.activation)
-
-        hy = permutate.permutate(hy, indices_array, axis=1, inv=True)
+        hys = [permutate.permutate(h, indices_array, axis=1, inv=True)
+               for h in result[:-1]]
+        trans_y = result[-1]
         ys = transpose_sequence.transpose_sequence(trans_y)
         ys = permutate_list(ys, indices, inv=True)
 
-        return hy, ys
+        return hys, ys
 
 
 class NStepRNNTanh(NStepRNNBase):
@@ -185,10 +217,11 @@ class NStepRNNTanh(NStepRNNBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
-        NStepRNNBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            use_bi_direction=False,  activation='tanh', **kwargs)
+    n_weights = 2
+    use_bi_direction = False
+
+    def rnn(self, *args):
+        return rnn.n_step_rnn(*args, activation='tanh')
 
 
 class NStepRNNReLU(NStepRNNBase):
@@ -223,10 +256,11 @@ class NStepRNNReLU(NStepRNNBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
-        NStepRNNBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            use_bi_direction=False, activation='relu', **kwargs)
+    n_weights = 2
+    use_bi_direction = False
+
+    def rnn(self, *args):
+        return rnn.n_step_rnn(*args, activation='relu')
 
 
 class NStepBiRNNTanh(NStepRNNBase):
@@ -262,10 +296,11 @@ class NStepBiRNNTanh(NStepRNNBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
-        NStepRNNBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            use_bi_direction=True, activation='tanh', **kwargs)
+    n_weights = 2
+    use_bi_direction = True
+
+    def rnn(self, *args):
+        return rnn.n_step_birnn(*args, activation='tanh')
 
 
 class NStepBiRNNReLU(NStepRNNBase):
@@ -300,7 +335,8 @@ class NStepBiRNNReLU(NStepRNNBase):
 
     """
 
-    def __init__(self, n_layers, in_size, out_size, dropout, **kwargs):
-        NStepRNNBase.__init__(
-            self, n_layers, in_size, out_size, dropout,
-            use_bi_direction=True, activation='relu', **kwargs)
+    n_weights = 2
+    use_bi_direction = True
+
+    def rnn(self, *args):
+        return rnn.n_step_birnn(*args, activation='relu')


Test patch: