Repository: chainer/chainer

Issue:
False warning is shown for BatchNormalization when batch size is 1, but tensor dimensions are larger than 2
https://github.com/chainer/chainer/blob/master/chainer/functions/normalization/batch_normalization.py#L90-L100

The above warning is correct only when the input tensor has 2 dimensions and batch size 1. E.g. with input `x = np.random.rand(1, 3, 4, 5)`, BatchNormalization works fine, since the statistics are aggregated over the 3rd and 4th axis as well, so this warning is confusing e.g. when training large semantic segmentation models that only fit in memory with batch size 1.


Discussion:

Resolving patch:
diff --git a/chainer/functions/normalization/batch_normalization.py b/chainer/functions/normalization/batch_normalization.py
--- a/chainer/functions/normalization/batch_normalization.py
+++ b/chainer/functions/normalization/batch_normalization.py
@@ -87,18 +87,6 @@ def forward(self, inputs):
         self.retain_inputs((0, 1))
         x, gamma, beta = inputs
 
-        if x.shape[0] == 1:
-            warnings.warn(
-                'A batch with no more than one sample has been given'
-                ' to F.batch_normalization. F.batch_normalization'
-                ' will always output a zero tensor for such batches.'
-                ' This could be caused by incorrect configuration in'
-                ' your code (such as running evaluation while'
-                ' chainer.config.train=True),'
-                ' but could also happen in the last batch of training'
-                ' if non-repeating iterator is used.',
-                UserWarning)
-
         xp = cuda.get_array_module(x)
         if self.running_mean is None:
             self.running_mean = xp.zeros_like(gamma)
@@ -107,6 +95,26 @@ def forward(self, inputs):
         self.axis = _compute_axis(x.ndim, gamma.ndim, self.axis)
         self.key_axis = _compute_key_axis(x.ndim, gamma.ndim, self.axis)
 
+        if all(x.shape[i] == 1 for i in self.axis):
+            if 0 in self.axis:
+                warnings.warn(
+                    'A batch with no more than one sample has been given'
+                    ' to F.batch_normalization. F.batch_normalization'
+                    ' will always output a zero tensor for such batches.'
+                    ' This could be caused by incorrect configuration in'
+                    ' your code (such as running evaluation while'
+                    ' chainer.config.train=True),'
+                    ' but could also happen in the last batch of training'
+                    ' if non-repeating iterator is used.',
+                    UserWarning)
+            else:
+                warnings.warn(
+                    'F.batch_normalization received a batch with single'
+                    ' dimensions along all axes that are used for aggregating'
+                    ' statistics. F.batch_normalization'
+                    ' will always output a zero tensor for such batches.',
+                    UserWarning)
+
         # TODO(niboshi): Refactor calculation of expander and axis into a
         # function and call it just before they are used.
 


Test patch:
diff --git a/tests/chainer_tests/functions_tests/normalization_tests/test_batch_normalization.py b/tests/chainer_tests/functions_tests/normalization_tests/test_batch_normalization.py
--- a/tests/chainer_tests/functions_tests/normalization_tests/test_batch_normalization.py
+++ b/tests/chainer_tests/functions_tests/normalization_tests/test_batch_normalization.py
@@ -2,6 +2,7 @@
 
 import numpy
 import six
+import warnings
 
 import chainer
 from chainer.backends import cuda
@@ -394,4 +395,39 @@ def test_invalid(self):
             functions.batch_normalization(*self.args, eps=2e-6)
 
 
+class TestBatchNormalizationWarning(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def create_batch(self, param_shape, x_shape):
+        dtype = numpy.float32
+        gamma = numpy.random.uniform(.5, 1, param_shape).astype(dtype)
+        beta = numpy.random.uniform(-1, 1, param_shape).astype(dtype)
+        x = numpy.random.uniform(-1, 1, x_shape).astype(dtype)
+        args = [x, gamma, beta]
+        return args
+
+    def test_invalid_batch(self):
+        args = self.create_batch((3,), (1, 3))
+        with testing.assert_warns(UserWarning):
+            functions.batch_normalization(*args)
+
+    def test_invalid_batch_no_batch_axis(self):
+        args = self.create_batch((1, 3,), (1, 3, 1))
+        with testing.assert_warns(UserWarning):
+            functions.batch_normalization(*args, axis=2)
+
+    def test_valid_batch(self):
+        args = self.create_batch((3,), (1, 3, 2, 2))
+        with warnings.catch_warnings(record=True) as w:
+            functions.batch_normalization(*args)
+            assert len(w) == 0
+
+    def test_valid_batch_no_batch_axis(self):
+        args = self.create_batch((1, 3,), (1, 3, 2))
+        with warnings.catch_warnings(record=True) as w:
+            functions.batch_normalization(*args, axis=2)
+            assert len(w) == 0
+
+
 testing.run_module(__name__, __file__)