Merge 4119eb0 into 6027634

chainerrl/links/mlp_bn.py

@@ -10,6 +10,8 @@
 from chainer import functions as F
 from chainer import links as L
 
+from chainerrl.initializers import LeCunNormal
+
 
 class LinearBN(chainer.Chain):
     """Linear layer with BatchNormalization."""
@@ -27,15 +29,28 @@ def __call__(self, x):
 
 
 class MLPBN(chainer.Chain):
-    """Multi-Layer Perceptron with BatchNormalization."""
+    """Multi-Layer Perceptron with Batch Normalization.
+
+    Args:
+        in_size (int): Input size.
+        out_size (int): Output size.
+        hidden_sizes (list of ints): Sizes of hidden channels.
+        normalize_input (bool): If set to True, Batch Normalization is applied
+            to inputs.
+        normalize_output (bool): If set to True, Batch Normalization is applied
+            to outputs.
+        nonlinearity (callable): Nonlinearity between layers.
+        last_wscale (float): Scale of weight initialization of the last layer.
+        """
 
     def __init__(self, in_size, out_size, hidden_sizes, normalize_input=True,
-                 normalize_output=False):
+                 normalize_output=False, nonlinearity=F.relu, last_wscale=1):
         self.in_size = in_size
         self.out_size = out_size
         self.hidden_sizes = hidden_sizes
         self.normalize_input = normalize_input
         self.normalize_output = normalize_output
+        self.nonlinearity = nonlinearity
 
         super().__init__()
         with self.init_scope():
@@ -49,9 +64,11 @@ def __init__(self, in_size, out_size, hidden_sizes, normalize_input=True,
                 for hin, hout in zip(hidden_sizes, hidden_sizes[1:]):
                     hidden_layers.append(LinearBN(hin, hout))
                 self.hidden_layers = chainer.ChainList(*hidden_layers)
-                self.output = L.Linear(hidden_sizes[-1], out_size)
+                self.output = L.Linear(hidden_sizes[-1], out_size,
+                                       initialW=LeCunNormal(last_wscale))
             else:
-                self.output = L.Linear(in_size, out_size)
+                self.output = L.Linear(in_size, out_size,
+                                       initialW=LeCunNormal(last_wscale))
 
             if normalize_output:
                 self.output_bn = L.BatchNormalization(out_size)
@@ -64,7 +81,7 @@ def __call__(self, x):
             h = self.input_bn(h)
         if self.hidden_sizes:
             for l in self.hidden_layers:
-                h = F.relu(l(h))
+                h = self.nonlinearity(l(h))
         h = self.output(h)
         if self.normalize_output:
             h = self.output_bn(h)

tests/links_tests/test_mlp_bn.py

@@ -0,0 +1,60 @@
+from __future__ import unicode_literals
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+from builtins import *  # NOQA
+from future import standard_library
+standard_library.install_aliases()
+
+import unittest
+
+import numpy as np
+
+import chainer
+import chainer.functions as F
+from chainer import testing
+from chainer.testing import attr
+
+import chainerrl
+
+
+@testing.parameterize(
+    *testing.product({
+        'in_size': [1, 5],
+        'out_size': [1, 3],
+        'hidden_sizes': [(), (1,), (1, 1), (7, 8)],
+        'normalize_input': [True, False],
+        'normalize_output': [True, False],
+        'nonlinearity': ['relu', 'elu'],
+        'last_wscale': [1, 1e-3],
+    })
+)
+class TestMLPBN(unittest.TestCase):
+
+    def _test_call(self, gpu):
+        nonlinearity = getattr(F, self.nonlinearity)
+        mlp = chainerrl.links.MLPBN(
+            in_size=self.in_size,
+            out_size=self.out_size,
+            hidden_sizes=self.hidden_sizes,
+            normalize_input=self.normalize_input,
+            normalize_output=self.normalize_output,
+            nonlinearity=nonlinearity,
+            last_wscale=self.last_wscale,
+        )
+        batch_size = 7
+        x = np.random.rand(batch_size, self.in_size).astype(np.float32)
+        if gpu >= 0:
+            mlp.to_gpu(gpu)
+            x = chainer.cuda.to_gpu(x)
+        y = mlp(x)
+        self.assertEqual(y.shape, (batch_size, self.out_size))
+        self.assertEqual(chainer.cuda.get_array_module(y),
+                         chainer.cuda.get_array_module(x))
+
+    def test_call_cpu(self):
+        self._test_call(gpu=-1)
+
+    @attr.gpu
+    def test_call_gpu(self):
+        self._test_call(gpu=0)