Repository: rwth-i6/returnn

Issue:
Fix batch norm default parameters
#508

New potential defaults:

* `momentum=0.1`
* `update_sample_only_in_training=True`
* `delay_sample_update=True`
* `param_version=2` (with #898)
* `masked_time` mandatory

See also other frameworks for defaults. Some references:

* [PyTorch-to-RETURNN comment](https://github.com/rwth-i6/pytorch-to-returnn/blob/0d55be95da12fbe3beaba8205a760aa5f18de2db/pytorch_to_returnn/torch/nn/modules/batchnorm.py#L90)
* [PyTorch BatchNorm1d](https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html)



Discussion:
One open question here is on our `masked_time` option, which is currently `True` by default. Which is the correct thing to do. This flattens the tensor, i.e. removes any padded frames, and calculates all the statistics only on the valid frames.

However, that option makes it slower and requires more memory. I don't really have recent numbers on this. Some benchmarks would be nice... This should also include back prop. Some old numbers are here: https://github.com/rwth-i6/returnn/pull/170#issuecomment-487600030 which indicate a factor of 2 in runtime.

Also, we currently do not use the fused TF op for batch norm, which would make it faster, again maybe factor of 2 in runtime. Although we could also use it on the flattened tensor. I don't know if there is any interplay between the two things.

Also note: All other frameworks I know do not use this padding, and always use a simple fused op. E.g. PyTorch (https://github.com/pytorch/pytorch/issues/70529) and Keras (https://github.com/tensorflow/tensorflow/issues/30161).

@vieting @JackTemaki @curufinwe some opinion would be nice.

Well, actually, the only question is, should we change that default? Or should we remove the default and make this a mandatory setting such that the user need to think about this?

It is clear that we should use the fused op in all cases, either on the flattened tensor or on the unflattened tensor. (There is the very old PR #170 on this, although we probably should redo this now.)

One note on the confusion of `momentum`: This probably comes from the different meaning of momentum.

RETURNN TF batch norm:
```
mean_new = (1 - momentum) * mean_old + momentum * mean_cur_batch
```

RETURNN Theano batch norm:
* cumulative moving average for mean. there is no momentum setting for exponential running average
* always the current batch statistics using the mean for variance (standard deviation), i.e. `T.sqrt(T.mean((x - sample_mean)**2,axis=0))`

[PyTorch batch norm](https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html):
```
mean_new = (1 - momentum) * mean_old + momentum * mean_cur_batch
```

[TensorFlow / Keras batch norm](https://keras.io/api/layers/normalization_layers/batch_normalization/):
```
mean_new = momentum * mean_old + (1 - momentum) * mean_cur_batch
```

[TensorFlow `tf.compat.v1.nn.fused_batch_norm`](https://www.tensorflow.org/api_docs/python/tf/compat/v1/nn/fused_batch_norm) (`exponential_avg_factor` is `momentum`):
```
mean_new = (1 - momentum) * mean_old + momentum * mean_cur_batch
```

[Flax batch norm](https://flax.readthedocs.io/en/latest/_modules/flax/linen/normalization.html#BatchNorm):
```
mean_new = momentum * mean_old + (1 - momentum) * mean_cur_batch
```

[MXNet batch norm](https://mxnet.apache.org/versions/1.6/api/r/docs/api/mx.symbol.BatchNorm.html)
```
mean_new = momentum * mean_old + (1 - momentum) * mean_cur_batch
```

[PaddlePaddle batch norm](https://www.paddlepaddle.org.cn/documentation/docs/en/1.8/api/dygraph/BatchNorm.html):
```
mean_new = momentum * mean_old + (1 - momentum) * mean_cur_batch
```

[cuDNN `cudnnBatchNormalizationForwardTraining`](https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnBatchNormalizationForwardTraining):
```
mean_new = (1 - momentum) * mean_old + momentum * mean_cur_batch
```

[PyTorch optimizer](https://pytorch.org/docs/stable/generated/torch.optim.SGD.html):
```
v_new = momentum * v_old + g
```

[Keras optimizer](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD):
```
v_new = momentum * v_old - learning_rate * g
```

One question is, which way do we want in returnn-common?

I also wonder whether we should couple it somehow to the learning rate or not (which would be quite non-standard).
(This is not really about the defaults though. In any case, the defaults should make sense, obviously.)

I also suggest `param_version=2` (via #898). All the automatic param conversion should work.

I also suggest to make `masked_time` mandatory (not having a default). While `masked_time=True` is maybe the correct thing to do, this can be much slower, and is also not consistent to how all other frameworks are doing it (like `masked_time=False`). So the user should explicitly decide on this.

Resolving patch:
diff --git a/returnn/tf/layers/base.py b/returnn/tf/layers/base.py
--- a/returnn/tf/layers/base.py
+++ b/returnn/tf/layers/base.py
@@ -1380,12 +1380,12 @@ def get_constraints_value(self):
 
   def batch_norm(self, data,
                  use_shift=True, use_std=True, use_sample=0.0, force_sample=False,
-                 momentum=0.99, epsilon=1e-3,
-                 update_sample_only_in_training=False,
-                 delay_sample_update=False,
-                 param_version=0,
+                 momentum=NotSpecified, epsilon=1e-3,
+                 update_sample_only_in_training=NotSpecified,
+                 delay_sample_update=NotSpecified,
+                 param_version=NotSpecified,
                  gamma_init=1.0, beta_init=0.0,
-                 masked_time=True):
+                 masked_time=NotSpecified):
     """
     :param Data data:
     :param bool use_shift:
@@ -1415,6 +1415,23 @@ def batch_norm(self, data,
       tf.nn.batch_normalization()
       https://github.com/deepmind/sonnet/blob/master/sonnet/python/modules/batch_norm.py
     """
+    from returnn.util import BehaviorVersion
+    # Note that the old defaults (behavior version <= 11) don't really make sense!
+    # https://github.com/rwth-i6/returnn/issues/522
+    if momentum is NotSpecified:
+      momentum = 0.1 if BehaviorVersion.get() >= 12 else 0.99
+    if update_sample_only_in_training is NotSpecified:
+      update_sample_only_in_training = True if BehaviorVersion.get() >= 12 else False
+    if delay_sample_update is NotSpecified:
+      delay_sample_update = True if BehaviorVersion.get() >= 12 else False
+    if param_version is NotSpecified:
+      param_version = 2 if BehaviorVersion.get() >= 12 else 0
+    BehaviorVersion.require(
+      masked_time is not NotSpecified, message="batch_norm masked_time should be specified explicitly", version=12)
+    if masked_time is NotSpecified:
+      assert BehaviorVersion.get() <= 11
+      masked_time = True
+
     with reuse_name_scope(self.get_absolute_name_scope_prefix() + "batch_norm", absolute=True):
       # Note about param_version:
       # We might later drop some earlier versions.
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -209,7 +209,7 @@ class BehaviorVersion:
   The version will be set after the config is defined at __main__.init_config() or Engine.__init__()
   """
 
-  _latest_behavior_version = 11
+  _latest_behavior_version = 12
   _behavior_version = None  # type: typing.Optional[int]
 
   @classmethod


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -554,7 +554,8 @@ def test_batch_norm_vars():
       "num_inputs": n_in,
       "network": {
         layer_name: {
-          "class": "linear", "activation": "relu", "batch_norm": True, "n_out": n_out, "is_output_layer": True,
+          "class": "linear", "activation": "relu", "batch_norm": {"masked_time": True},
+          "n_out": n_out, "is_output_layer": True,
           "from": "data:data"}
       }})
     network = TFNetwork(config=config, train_flag=True)
@@ -564,15 +565,15 @@ def test_batch_norm_vars():
     print("layer vars:")
     pprint(layer.params)
     assert layer.use_batch_norm
-    bn_prefix = "batch_norm/%s_%s_output_" % (layer_name, layer_name)
+    bn_prefix = "batch_norm/v2_"
     assert_equal(set(layer.params.keys()), {
       "W", "b", bn_prefix + "beta", bn_prefix + "mean", bn_prefix + "gamma", bn_prefix + "variance"})
     assert_equal(layer.params["W"].get_shape().as_list(), [n_in, n_out])
     assert_equal(layer.params["b"].get_shape().as_list(), [n_out])
-    assert_equal(layer.params[bn_prefix + "beta"].get_shape().as_list(), [1, 1, n_out])
-    assert_equal(layer.params[bn_prefix + "gamma"].get_shape().as_list(), [1, 1, n_out])
-    assert_equal(layer.params[bn_prefix + "mean"].get_shape().as_list(), [1, 1, n_out])
-    assert_equal(layer.params[bn_prefix + "variance"].get_shape().as_list(), [1, 1, n_out])
+    assert_equal(layer.params[bn_prefix + "beta"].get_shape().as_list(), [n_out])
+    assert_equal(layer.params[bn_prefix + "gamma"].get_shape().as_list(), [n_out])
+    assert_equal(layer.params[bn_prefix + "mean"].get_shape().as_list(), [n_out])
+    assert_equal(layer.params[bn_prefix + "variance"].get_shape().as_list(), [n_out])
 
 
 def _test_batch_norm_param_old_to_new_import(old_version, new_version):
@@ -586,7 +587,7 @@ def _make_net_dict(param_version):
     return {
       layer_name: {
         "class": "batch_norm", "from": "data:data", "is_output_layer": True,
-        "param_version": param_version}
+        "param_version": param_version, "masked_time": True}
     }
 
   with make_scope() as session:
@@ -674,7 +675,7 @@ def test_batch_norm():
       feed_dict={
         src_nhwc.output.placeholder: input_data,
         src_nhwc.output.size_placeholder[0]: seq_lens})
-    assert numpy.array_equal(out_1, out_2)
+    numpy.testing.assert_array_almost_equal(out_1, out_2)
     print(numpy.sum(out_1 - out_2))
 
 
@@ -903,14 +904,14 @@ def test_cnn_building_block():
                "in_spatial_dims": ("T", "dim:6"),
                "strides": (1, 1), "dilation_rate": (1, 1), "padding": "SAME", "activation": None, "with_bias": False,
                "from": "swap_axes"},
-        "bn1": {"class": "batch_norm", "from": "c1"},
+        "bn1": {"class": "batch_norm", "from": "c1", "masked_time": True},
         "y1": {"class": "activation", "activation": "relu", "batch_norm": False, "from": "bn1"},
         "c2": {"class": "conv", "n_out": filters, "filter_size": filter_size, "auto_use_channel_first": False,
                "in_spatial_dims": ("T", "dim:6"),
                "strides": (1, 1), "dilation_rate": (1, 1), "padding": "SAME", "activation": None, "with_bias": False,
                "from": "y1"},
         "p": {"class": "combine", "kind": "add", "from": ["c2", "swap_axes"]},
-        "bn2": {"class": "batch_norm", "from": "p"},
+        "bn2": {"class": "batch_norm", "from": "p", "masked_time": True},
         "y2": {"class": "activation", "activation": "relu", "batch_norm": False, "from": "bn2"},
 
         "out_pool": {