Repository: chainer/chainer

Issue:
Defer initialization of size of BatchNormalization and infer it at first iteration.
The input dimension of several links (e.g. `L.Linear`, `L.Convolution`) do not have to be fed and is inferred by the minibatch size of the first iteration. Theoretically we should do the similar thing with `L.BatchNormalization`.


Discussion:

Resolving patch:
diff --git a/chainer/links/normalization/batch_normalization.py b/chainer/links/normalization/batch_normalization.py
--- a/chainer/links/normalization/batch_normalization.py
+++ b/chainer/links/normalization/batch_normalization.py
@@ -33,8 +33,9 @@ class BatchNormalization(link.Link):
     fine-tuning mode.
 
     Args:
-        size (int or tuple of ints): Size (or shape) of channel
-            dimensions.
+        size (int, tuple of ints, or None): Size (or shape) of channel
+            dimensions.  If ``None``, the size will be determined from
+            dimension(s) of the input batch during the first forward pass.
         decay (float): Decay rate of moving average. It is used on training.
         eps (float): Epsilon value for numerical stability.
         dtype (numpy.dtype): Type to use in computing.
@@ -68,35 +69,143 @@ class BatchNormalization(link.Link):
         ~BatchNormalization.eps (float): Epsilon value for numerical stability.
             This value is added to the batch variances.
 
+    .. admonition:: Example
+
+        >>> x = np.arange(12).reshape(4, 3).astype(np.float32) ** 2
+        >>> x
+        array([[  0.,   1.,   4.],
+               [  9.,  16.,  25.],
+               [ 36.,  49.,  64.],
+               [ 81., 100., 121.]], dtype=float32)
+        >>> bn = chainer.links.BatchNormalization(3)
+        >>> bn(x)
+        variable([[-1.        , -1.0664359 , -1.1117983 ],
+                  [-0.71428573, -0.6714596 , -0.6401263 ],
+                  [ 0.14285715,  0.19748813,  0.23583598],
+                  [ 1.5714287 ,  1.5404074 ,  1.5160885 ]])
+        >>> (x - x.mean(axis=0)) / np.sqrt(x.var(axis=0) + 2e-5)
+        array([[-1.        , -1.0664359 , -1.1117983 ],
+               [-0.71428573, -0.6714596 , -0.6401263 ],
+               [ 0.14285715,  0.19748813,  0.235836  ],
+               [ 1.5714285 ,  1.5404074 ,  1.5160886 ]], dtype=float32)
+
+        There are several ways to make a BatchNormalization link.
+        Consider an input of batched 10 images of 32x32 with 3 channels.
+
+        >>> x = np.random.randn(10, 3, 32, 32).astype(np.float32)
+
+        1. Give the parameter size:
+
+            To normalize for each channel, give the number of channels
+            to ``size``.
+
+            >>> bn = chainer.links.BatchNormalization(3)
+            >>> bn.avg_mean.shape
+            (3,)
+            >>> bn.beta += 2.0
+            >>> bn.gamma *= 5.0
+            >>> list(sorted(bn.namedparams()))  # doctest: +ELLIPSIS
+            [('/beta', variable([2., ...])), ('/gamma', variable([5., ...]))]
+            >>> y = bn(x)
+            >>> y.shape
+            (10, 3, 32, 32)
+            >>> np.testing.assert_allclose(
+            ...     y.array.mean(axis=(0, 2, 3)), bn.beta.array, atol=1e-6)
+            >>> np.testing.assert_allclose(
+            ...     y.array.std(axis=(0, 2, 3)),
+            ...     bn.gamma.array, atol=1e-3)
+
+            To normalize for each channel for each pixel, ``size`` should
+            be the tuple of the dimensions.
+
+            >>> bn = chainer.links.BatchNormalization((3, 32, 32))
+            >>> bn.avg_mean.shape
+            (3, 32, 32)
+            >>> y = bn(x)
+            >>> y.shape
+            (10, 3, 32, 32)
+            >>> np.testing.assert_allclose(
+            ...     y.array.mean(axis=0), bn.beta.array, atol=1e-6)
+            >>> np.testing.assert_allclose(
+            ...     y.array.std(axis=0),
+            ...     bn.gamma.array, atol=1e-3)
+
+            By default, channel axis is (or starts from) the 1st axis of the
+            input shape.
+
+        2. Give the aggregate axes:
+
+            from Chainer v5
+
+            With ``axis`` option, similarly to NumPy, you may specify the
+            aggregate axes, which are treated as the "batch" axes for the
+            batch statistics.
+
+            You can omit ``size`` if ``axis`` is given. In this case, creation
+            of persistent values ``avg_mean``, ``avg_var`` and parameters
+            ``beta``, ``gamma`` is deferred until first forward propagation.
+
+            The examples in 1. corresponds to the following, respectively.
+
+            >>> bn = chainer.links.BatchNormalization(axis=(0, 2, 3))
+            >>> hasattr(bn, 'avg_mean')
+            False
+            >>> y = bn(x)
+            >>> bn.avg_mean.shape
+            (3,)
+
+            >>> bn = chainer.links.BatchNormalization(axis=0)
+            >>> hasattr(bn, 'avg_mean')
+            False
+            >>> y = bn(x)
+            >>> bn.avg_mean.shape
+            (3, 32, 32)
+
     """
 
-    def __init__(self, size, decay=0.9, eps=2e-5, dtype=numpy.float32,
+    def __init__(self, size=None, decay=0.9, eps=2e-5, dtype=numpy.float32,
                  use_gamma=True, use_beta=True,
                  initial_gamma=None, initial_beta=None, axis=None):
         super(BatchNormalization, self).__init__()
-        self.avg_mean = numpy.zeros(size, dtype=dtype)
-        self.register_persistent('avg_mean')
-        self.avg_var = numpy.zeros(size, dtype=dtype)
-        self.register_persistent('avg_var')
+
+        if size is None and axis is None:
+            raise RuntimeError('size or axis is required')
         self.N = 0
         self.register_persistent('N')
         self.decay = decay
         self.eps = eps
+        if isinstance(axis, int):
+            axis = (axis,)
         self.axis = axis
+        self._dtype = dtype
 
         with self.init_scope():
             if use_gamma:
                 if initial_gamma is None:
                     initial_gamma = 1
-                initial_gamma = initializers._get_initializer(initial_gamma)
-                initial_gamma.dtype = dtype
-                self.gamma = variable.Parameter(initial_gamma, size)
+                gamma_initializer = \
+                    initializers._get_initializer(initial_gamma)
+                gamma_initializer.dtype = self._dtype
+                self.gamma = variable.Parameter(gamma_initializer)
             if use_beta:
                 if initial_beta is None:
                     initial_beta = 0
-                initial_beta = initializers._get_initializer(initial_beta)
-                initial_beta.dtype = dtype
-                self.beta = variable.Parameter(initial_beta, size)
+                beta_initializer = initializers._get_initializer(initial_beta)
+                beta_initializer.dtype = self._dtype
+                self.beta = variable.Parameter(beta_initializer)
+
+        if size is not None:
+            self._initialize_params(size)
+
+    def _initialize_params(self, shape):
+        self.avg_mean = numpy.zeros(shape, dtype=self._dtype)
+        self.register_persistent('avg_mean')
+        self.avg_var = numpy.zeros(shape, dtype=self._dtype)
+        self.register_persistent('avg_var')
+        if hasattr(self, 'gamma'):
+            self.gamma.initialize(shape)
+        if hasattr(self, 'beta'):
+            self.beta.initialize(shape)
 
     def __call__(self, x, **kwargs):
         """__call__(self, x, finetune=False)
@@ -127,6 +236,13 @@ def __call__(self, x, **kwargs):
             test='test argument is not supported anymore. '
                  'Use chainer.using_config')
 
+        if not hasattr(self, 'avg_mean'):
+            param_shape = tuple([
+                d
+                for i, d in enumerate(x.shape)
+                if i not in self.axis])
+            self._initialize_params(param_shape)
+
         if hasattr(self, 'gamma'):
             gamma = self.gamma
         else:


Test patch:
diff --git a/tests/chainer_tests/links_tests/normalization_tests/test_batch_normalization.py b/tests/chainer_tests/links_tests/normalization_tests/test_batch_normalization.py
--- a/tests/chainer_tests/links_tests/normalization_tests/test_batch_normalization.py
+++ b/tests/chainer_tests/links_tests/normalization_tests/test_batch_normalization.py
@@ -409,4 +409,56 @@ def test_redundant_argument(self):
             self.link(self.x, unknown_argument=1)
 
 
+@testing.parameterize(
+    {'shape': (5, 4, 3, 2), 'axis': (0, 2, 3)},
+    {'shape': (5, 4), 'axis': 0},
+    {'shape': (5, 4, 3), 'axis': (0, 1)},
+)
+class TestChannalSizeInference(unittest.TestCase):
+
+    def setUp(self):
+        self.x = numpy.random.randn(*self.shape).astype('f')
+
+        axis = self.axis
+        if isinstance(axis, int):
+            axis = (axis,)
+        self.expected_size = tuple(
+            n
+            for i, n in enumerate(self.shape)
+            if i not in axis
+        )
+
+    def test_no_inference(self):
+        bn = links.BatchNormalization(self.expected_size)
+        assert hasattr(bn, 'avg_mean')
+        assert hasattr(bn, 'avg_var')
+
+    def test_inference(self):
+        bn = links.BatchNormalization(axis=self.axis)
+        bn(self.x)
+        assert bn.beta.shape == self.expected_size
+        assert bn.gamma.shape == self.expected_size
+        assert bn.avg_mean.shape == self.expected_size
+        assert bn.avg_var.shape == self.expected_size
+
+    def test_no_gamma(self):
+        bn = links.BatchNormalization(axis=self.axis, use_gamma=False)
+        assert not hasattr(bn, 'gamma')
+        bn(self.x)
+        assert not hasattr(bn, 'gamma')
+
+    def test_no_beta(self):
+        bn = links.BatchNormalization(axis=self.axis, use_beta=False)
+        assert not hasattr(bn, 'beta')
+        bn(self.x)
+        assert not hasattr(bn, 'beta')
+
+
+class TestFailChannalSizeInference(unittest.TestCase):
+
+    def test_fail_inference(self):
+        with self.assertRaises(RuntimeError):
+            links.BatchNormalization()
+
+
 testing.run_module(__name__, __file__)