Repository: chainer/chainer

Issue:
Serialization of BatchNormalization statistics fails when layer is initialized using batch axis
BatchNormalization serialization fails to load the stored statistics `avg_mean` and `avg_var` when the layer is initialized using the `axis` parameter rather than `size` parameter. As a result, at test time the BatchNorm layer uses default initialization of 0 and 1 for mean and variance rather than the saved statistics. This is because the `avg_mean` and `avg_var` attributes are initialized and registered as persistent in the `_initialize_params()` method ([link](https://github.com/chainer/chainer/blob/master/chainer/links/normalization/batch_normalization.py#L238)) which is only called on the first forward pass, since the feature dimensions must be inferred from the data. As a result they are skipped by the serializer when loading parameters into a new test-time model prior to passing any data through the model.

To be fixed, the `avg_mean` and `avg_var` attributes would need to be registered as persistent in the BatchNormalization `__init__()` method. In addition, in the loop over persistent attributes in chainer.Link located [here](https://github.com/chainer/chainer/blob/master/chainer/link.py#L663), there would need to be a case handling unitialized persistent attributes that can read the dimension from the serializer in a similar manner to how it is done for parameters.

Here is an example script that demonstrates the problem:
```
import chainer
import numpy as np
from chainer.links import BatchNormalization
from chainer.serializers import load_npz, save_npz

batch_sample = np.random.rand(5, 4, 3, 2).astype(np.float32)

norm_1 = BatchNormalization(axis=(0, 1))

print('norm_1 pre-call persistent values:', norm_1._persistent)
print('norm_1 avg_mean:')
print(norm_1.avg_mean)

norm_1(batch_sample)

print('norm_1 post-call persistent values:', norm_1._persistent)
print('norm_1 avg_mean:')
print(norm_1.avg_mean)
print('')

save_npz('norm.npz', norm_1)

npz = np.load('norm.npz')

print('norm.npz files:', npz.files)
print('norm.npz avg_mean:')
print(npz['avg_mean'])
print('')

chainer.config.train = False

norm_2 = BatchNormalization(axis=(0, 1))
load_npz('norm.npz', norm_2)

print('norm_2 pre-call persistent values:', norm_2._persistent)
print('norm_2 avg_mean:')
print(norm_2.avg_mean)

norm_2(batch_sample)

print('norm_2 post-call persistent values:', norm_2._persistent)
print('norm_2 avg_mean:')
print(norm_2.avg_mean)
```
along with its output:
```
norm_1 pre-call persistent values: {'N'}
norm_1 avg_mean:
None
norm_1 post-call persistent values: {'N', 'avg_mean', 'avg_var'}
norm_1 avg_mean:
[[0.04254831 0.0445514 ]
 [0.04263044 0.04522393]
 [0.05117331 0.05540558]]

norm.npz files: ['gamma', 'beta', 'N', 'avg_mean', 'avg_var']
norm.npz avg_mean:
[[0.04254831 0.0445514 ]
 [0.04263044 0.04522393]
 [0.05117331 0.05540558]]

norm_2 pre-call persistent values: {'N'}
norm_2 avg_mean:
None
norm_2 post-call persistent values: {'N', 'avg_mean', 'avg_var'}
norm_2 avg_mean:
[[0. 0.]
 [0. 0.]
 [0. 0.]]
```
We can see that the avg_mean attribute only exists after a call, and thus even though it is correctly present in the npz archive, those attributes only get created after the parameters have been loaded and use the default initialization values of 0.

And finally, in case it matters, here's the output of print_runtime_info().
```
>>> chainer.print_runtime_info()
Platform: Linux-4.4.0-72-generic-x86_64-with-debian-stretch-sid
Chainer: 6.5.0
ChainerX: Not Available
NumPy: 1.17.1
CuPy:
  CuPy Version          : 6.5.0
  CUDA Root             : /cm/shared/apps/cuda10.0/toolkit/10.0.130
  CUDA Build Version    : 10000
  CUDA Driver Version   : 10010
  CUDA Runtime Version  : 10000
  cuDNN Build Version   : 7603
  cuDNN Version         : 7603
  NCCL Build Version    : 2402
  NCCL Runtime Version  : 2402
iDeep: Not Available
```


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
@@ -190,8 +190,6 @@ class BatchNormalization(link.Link):
 
     gamma = None
     beta = None
-    avg_mean = None
-    avg_var = None
 
     def __init__(self, size=None, decay=0.9, eps=2e-5, dtype=None,
                  use_gamma=True, use_beta=True,
@@ -229,16 +227,19 @@ def __init__(self, size=None, decay=0.9, eps=2e-5, dtype=None,
                 beta_initializer.dtype = self._highprec_dtype
                 self.beta = variable.Parameter(beta_initializer)
 
-        if size is not None:
+        if size is None:
+            self.avg_mean = None
+            self.avg_var = None
+        else:
             self._initialize_params(size)
+        self.register_persistent('avg_mean')
+        self.register_persistent('avg_var')
 
     def _initialize_params(self, shape):
         self.avg_mean = self._init_array(self._initial_avg_mean, 0, shape)
         self._initial_avg_mean = None
-        self.register_persistent('avg_mean')
         self.avg_var = self._init_array(self._initial_avg_var, 1, shape)
         self._initial_avg_var = None
-        self.register_persistent('avg_var')
         if self.gamma is not None:
             self.gamma.initialize(shape)
         if self.beta is not None:


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
@@ -656,4 +656,48 @@ def test_lazy_initialization_with_non_zero_current_cuda_device(self):
         assert backend.GpuDevice.from_array(bn.avg_var) == device
 
 
+@testing.parameterize(*testing.product({
+    'x_shape,bn_kwargs': [
+        ((4, 3), {'axis': (0,)}),
+        ((4, 3), {'size': (3,)}),
+    ],
+}))
+class TestSerialize(unittest.TestCase):
+
+    def create_link(self):
+        return links.BatchNormalization(**self.bn_kwargs)
+
+    def train_link(self, bn):
+        x = numpy.random.rand(*self.x_shape).astype(numpy.float32)
+        bn(x)
+        x = numpy.random.rand(*self.x_shape).astype(numpy.float32)
+        bn(x, finetune=True)
+        # has non-trivial values to be stored
+        assert bn.avg_mean is not None
+        assert bn.N == 1
+
+    def create_serializer_pair(self):
+        target = {}
+        return (
+            chainer.serializers.DictionarySerializer(target),
+            chainer.serializers.NpzDeserializer(target),
+        )
+
+    def test_serialize(self):
+        ser, de = self.create_serializer_pair()
+
+        link1 = self.create_link()
+        self.train_link(link1)
+        link1.serialize(ser)
+
+        link2 = self.create_link()
+        link2.serialize(de)
+
+        testing.assert_allclose(link2.avg_mean, link1.avg_mean)
+        testing.assert_allclose(link2.avg_var, link1.avg_var)
+        testing.assert_allclose(link2.beta.array, link1.beta.array)
+        testing.assert_allclose(link2.gamma.array, link1.gamma.array)
+        assert link2.N == link1.N
+
+
 testing.run_module(__name__, __file__)