returnn 1.20230504.150134__tar.gz → 1.20230504.233443__tar.gz

{returnn-1.20230504.150134/returnn.egg-info → returnn-1.20230504.233443}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230504.150134
+Version: 1.20230504.233443
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20230504.233443/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230504.233443'
+long_version = '1.20230504.233443+git.e9a9920'

{returnn-1.20230504.150134 → returnn-1.20230504.233443}/returnn/frontend/math_.py

@@ -4,7 +4,7 @@ Math ops
 
 from __future__ import annotations
 import typing
-from typing import Optional, Sequence, Union, Tuple
+from typing import Optional, Sequence, Union, Tuple, overload
 import numpy
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
@@ -33,6 +33,7 @@ __all__ = [
     "logical_and",
     "logical_or",
     "logical_not",
+    "opt_logical_or",
     "maximum",
     "minimum",
     "identity",
@@ -299,6 +300,24 @@ def logical_not(a: Tensor) -> Tensor:
     return a._raw_backend.activation(a, "logical_not")
 
 
+@overload
+def opt_logical_or(a: bool, b: bool) -> bool:
+    """logical or"""
+
+
+def opt_logical_or(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tensor, bool]:
+    """logical or"""
+    if isinstance(a, bool):
+        if a:
+            return True
+        return b
+    if isinstance(b, bool):
+        if b:
+            return True
+        return a
+    return combine(a, "logical_or", b)
+
+
 def maximum(a: Tensor, b: Union[Tensor, _RawTensorTypes], *other_tensors) -> Tensor:
     """maximum"""
     if not other_tensors:

{returnn-1.20230504.150134 → returnn-1.20230504.233443}/returnn/frontend/normalization.py

@@ -47,8 +47,7 @@ class LayerNorm(rf.Module):
         self.bias.initial = 0.0
 
     def __call__(self, x: Tensor) -> Tensor:
-        mean = rf.reduce_mean(x, axis=self.in_dim)
-        variance = rf.reduce_mean(rf.squared_difference(x, mean), axis=self.in_dim)
+        mean, variance = rf.moments(x, axis=self.in_dim)
         norm_x = (x - mean) * rf.rsqrt(variance + self.eps)
         return norm_x * self.scale + self.bias
 
@@ -93,14 +92,14 @@ class BatchNorm(rf.Module):
         *,
         affine: bool = True,
         momentum: float = 0.1,
-        epsilon: float = 1e-3,
+        eps: float = 1e-3,
         use_mask: Optional[bool] = None,
     ):
         """
         :param in_dim: the feature dimension of the input
         :param affine: whether to use learnable parameters gamma and beta
         :param momentum: momentum for the running mean and variance
-        :param epsilon: epsilon for the variance
+        :param eps: epsilon for the variance
         :param use_mask: whether to use a mask for dynamic spatial dims.
           This must be specified if the input has dynamic spatial dims.
           True would use the correct masking then. However, that is inconsistent to all other frameworks
@@ -113,7 +112,7 @@ class BatchNorm(rf.Module):
         self.in_dim = in_dim
         self.use_mask = use_mask
         self.momentum = momentum
-        self.epsilon = epsilon
+        self.eps = eps
         self.running_mean = rf.Parameter([in_dim], auxiliary=True)
         self.running_mean.initial = 0.0
         self.running_variance = rf.Parameter([in_dim], auxiliary=True)
@@ -128,11 +127,8 @@ class BatchNorm(rf.Module):
             self.beta.initial = 0.0
 
     def __call__(self, source: Tensor) -> Tensor:
-        # We wrap the RETURNN layer because we want efficient handling if possible,
-        # which is potentially the use of a fused op,
-        # and maybe reordering of dims.
-        # https://github.com/rwth-i6/returnn_common/issues/89
         assert self.in_dim in source.dims
+
         if any(d.need_masking() for d in source.dims if d != self.in_dim):
             if self.use_mask is None:
                 raise ValueError(
@@ -141,6 +137,44 @@ class BatchNorm(rf.Module):
             use_mask = self.use_mask
         else:
             use_mask = False  # not needed. False because this potentially enables an efficient fused op.
+
+        if use_mask:
+            # Generic implementation which supports masking.
+            use_current_batch_stats = self.running_mean is None or rf.get_run_ctx().train_flag
+            update_running_stats = self.running_mean is not None and rf.get_run_ctx().train_flag
+            need_current_batch_stats = rf.opt_logical_or(use_current_batch_stats, update_running_stats)
+
+            mean_cur_batch, variance_cur_batch = rf.cond(
+                need_current_batch_stats,
+                lambda: rf.moments(source, axis=[d for d in source.dims if d != self.in_dim]),
+                lambda: (self.running_mean, self.running_variance),
+            )
+
+            def _update_running_stats():
+                self.running_mean.assign_add((mean_cur_batch - self.running_mean) * self.momentum)
+                self.running_variance.assign_add((variance_cur_batch - self.running_variance) * self.momentum)
+
+            rf.cond(update_running_stats, _update_running_stats, lambda: None)
+
+            mean, variance = rf.cond(
+                use_current_batch_stats,
+                lambda: (mean_cur_batch, variance_cur_batch),
+                lambda: (self.running_mean, self.running_variance),
+            )
+
+            bn = (source - mean) * rf.rsqrt(variance + self.eps)
+            if self.gamma is not None:
+                bn *= self.gamma
+            if self.beta is not None:
+                bn += self.beta
+            return bn
+
+        # Fallback to specific backend implementation for the standard case without masking.
+        # This fallback probably can internally use a more efficient implementation like from CuDNN.
+        # In case of TF-net-dict backend, we wrap the RETURNN layer because we want efficient handling if possible,
+        # which is potentially the use of a fused op,
+        # and maybe reordering of dims.
+        # https://github.com/rwth-i6/returnn_common/issues/89
         # noinspection PyProtectedMember
         return source._raw_backend.batch_norm(
             source=source,
@@ -148,7 +182,7 @@ class BatchNorm(rf.Module):
             use_mask=use_mask,
             affine=self.affine,
             momentum=self.momentum,
-            epsilon=self.epsilon,
+            epsilon=self.eps,
             running_mean=self.running_mean,
             running_variance=self.running_variance,
             gamma=self.gamma,

{returnn-1.20230504.150134 → returnn-1.20230504.233443}/returnn/torch/updater.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 
 import torch
 import typing
-from typing import Set
+from typing import Any, Set, Dict
 
 from returnn.log import log
 from returnn.util.basic import RefIdEq
@@ -35,7 +35,8 @@ def _init_optimizer_classes_dict():
 
 def get_optimizer_class(class_name):
     """
-    :param str|function|type[torch.optim.Optimizer] class_name: Optimizer data, e.g. "adam", torch.optim.Adam...
+    :param str|()->torch.optim.Optimizer|type[torch.optim.Optimizer] class_name:
+        Optimizer data, e.g. "adam", torch.optim.Adam...
     :return: Optimizer class
     :rtype: type[torch.optim.Optimizer]
     """
@@ -156,7 +157,7 @@ class Updater(object):
         if isinstance(optimizer_opts, torch.optim.Optimizer):
             return optimizer_opts
         elif callable(optimizer_opts):
-            optimizer_opts = {"class": optimizer_opts}
+            optimizer_opts: Dict[str, Any] = {"class": optimizer_opts}
         else:
             if not isinstance(optimizer_opts, dict):
                 raise ValueError("'optimizer' must of type dict, callable or torch.optim.Optimizer instance.")
@@ -246,18 +247,22 @@ class Updater(object):
         # Tracker of visited parameters to only add each parameter once, in case two modules share common parameters.
         # We need the wrapper class RefIdEq because Parameters are compared by value and not by reference.
         visited_params: Set[RefIdEq[torch.nn.Parameter]] = set()
-        for mn, m in self.network.named_modules():
-            for pn, p in m.named_parameters():
-                if RefIdEq(p) in visited_params:
+        for module_name, module in self.network.named_modules():
+            module_name: str
+            module: torch.nn.Module
+            for param_name, param in module.named_parameters(recurse=False):
+                param_name: str
+                param: torch.nn.Parameter
+                if RefIdEq(param) in visited_params:
                     continue
-                visited_params.add(RefIdEq(p))
-                fpn = "%s.%s" % (mn, pn) if mn else pn  # Full param name
-                if pn.endswith("bias"):
-                    no_wd_params.add(fpn)
-                elif pn.endswith("weight") and isinstance(m, blacklist_wd_modules):
-                    no_wd_params.add(fpn)
+                visited_params.add(RefIdEq(param))
+                full_param_name = "%s.%s" % (module_name, param_name) if module_name else param_name
+                if param_name.endswith("bias"):
+                    no_wd_params.add(full_param_name)
+                elif param_name.endswith("weight") and isinstance(module, blacklist_wd_modules):
+                    no_wd_params.add(full_param_name)
                 else:
-                    wd_params.add(fpn)
+                    wd_params.add(full_param_name)
 
         param_dict = {pn: p for pn, p in self.network.named_parameters()}
         optim_groups = [

{returnn-1.20230504.150134 → returnn-1.20230504.233443/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230504.150134
+Version: 1.20230504.233443
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20230504.150134 → returnn-1.20230504.233443}/tests/test_rf_normalization.py

@@ -37,3 +37,29 @@ def test_batch_norm():
         out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
 
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
+
+
+def test_batch_norm_masking():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+        }
+    )
+
+    class _Net(rf.Module):
+        def __init__(self):
+            super().__init__()
+            self.bn = rf.BatchNorm(in_dim, use_mask=True)
+
+        def __call__(self, out: Tensor) -> Tensor:
+            out = self.bn(out)
+            return out
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: _Net, extern_data: TensorDict):
+        out = model(extern_data["data"])
+        out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)

returnn-1.20230504.150134/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230504.150134'
-long_version = '1.20230504.150134+git.3cad0e0'