returnn 1.20240617.171744__tar.gz → 1.20240618.151147__tar.gz

{returnn-1.20240617.171744/returnn.egg-info → returnn-1.20240618.151147}/PKG-INFO

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

returnn-1.20240618.151147/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240618.151147'
+long_version = '1.20240618.151147+git.84cb714'

{returnn-1.20240617.171744 → returnn-1.20240618.151147}/returnn/frontend/__init__.py

@@ -21,6 +21,7 @@ from .state import *
 from .array_ import *
 from .attention import *
 from .backend import *
+from .build_from_dict import *
 from .cond import *
 from .const import *
 from .container import *
@@ -46,6 +47,7 @@ from .rec import *
 from .reduce import *
 from .run_ctx import *
 from .signal import *
+from .stepwise_scheduler import *
 from .types import *
 
 # Modules not in the main namespace but in sub namespaces.

returnn-1.20240618.151147/returnn/frontend/build_from_dict.py

@@ -0,0 +1,60 @@
+"""
+Construct modules (or other objects) from dictionaries.
+"""
+
+from __future__ import annotations
+from typing import Union, Any, Type, Dict
+import importlib
+import returnn.frontend as rf
+
+
+__all__ = ["build_from_dict", "build_dict"]
+
+
+def build_from_dict(d: Dict[str, Any], *args, **kwargs) -> Union[rf.Module, Any]:
+    """
+    Build a module (or other object) from a dictionary.
+    `"class"` in the dict is required and specifies the class to be instantiated.
+    The other options are passed to the class constructor.
+
+    :param d: dictionary with the class name and other options
+    :param args: passed to the class constructor
+    :param kwargs: passed to the class constructor
+    :return: cls(*args, **d, **kwargs)
+    """
+    if "class" not in d:
+        raise ValueError(f"build_from_dict: Missing 'class' key in dict: {d}")
+    d = d.copy()
+    cls_name = d.pop("class")
+    cls = _get_cls(cls_name)
+    return cls(*args, **d, **kwargs)
+
+
+def build_dict(cls: Type, **kwargs) -> Dict[str, Any]:
+    """
+    Build a dictionary for :func:`build_from_dict`.
+    The class name is stored in the `"class"` key.
+
+    Note that this is intended to be used for serialization
+    and also to get a unique stable hashable representation
+    (e.g. for Sisyphus :func:`sis_hash_helper`)
+    which should not change if the class is renamed or moved
+    to keep the hash stable.
+    """
+    return {"class": _get_cls_name(cls), **kwargs}
+
+
+def _get_cls(cls_name: str) -> Type:
+    if "." not in cls_name:
+        raise ValueError(f"Expected '.' in class name: {cls_name}")
+    mod_name, cls_name = cls_name.rsplit(".", 1)
+    if mod_name == "rf":
+        return getattr(rf, cls_name)
+    mod = importlib.import_module(mod_name)
+    return getattr(mod, cls_name)
+
+
+def _get_cls_name(cls: Type) -> str:
+    if getattr(rf, cls.__name__, None) is cls:
+        return f"rf.{cls.__name__}"
+    return f"{cls.__module__}.{cls.__name__}"

{returnn-1.20240617.171744 → returnn-1.20240618.151147}/returnn/frontend/encoder/conformer.py

@@ -227,6 +227,10 @@ class ConformerEncoderLayer(rf.Module):
             conv_norm = rf.BatchNorm(out_dim, **conv_norm_opts)
         elif isinstance(conv_norm, type):
             conv_norm = conv_norm(out_dim, **(conv_norm_opts or {}))
+        elif isinstance(conv_norm, dict):
+            conv_norm = rf.build_from_dict(conv_norm, out_dim, **(conv_norm_opts or {}))
+        if not callable(conv_norm):
+            raise TypeError(f"{self}: unexpected conv_norm type {conv_norm!r}")
         self.conv_block = ConformerConvBlock(out_dim=out_dim, kernel_size=conv_kernel_size, norm=conv_norm)
         self.conv_layer_norm = rf.LayerNorm(out_dim)
 

{returnn-1.20240617.171744 → returnn-1.20240618.151147}/returnn/frontend/normalization.py

@@ -8,7 +8,7 @@ from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 
 
-__all__ = ["moments", "LayerNorm", "BatchNorm", "normalize", "Normalize"]
+__all__ = ["moments", "LayerNorm", "GroupNorm", "BatchNorm", "normalize", "Normalize"]
 
 
 def moments(
@@ -70,6 +70,30 @@ class LayerNorm(rf.Module):
         return norm_x * self.scale + self.bias
 
 
+class GroupNorm(rf.Module):
+    """
+    `Group normalization <https://arxiv.org/abs/1803.08494>`__.
+    """
+
+    def __init__(self, in_dim: Union[rf.Dim, Sequence[rf.Dim]], *, num_groups: Union[int, Dim], eps: float = 1e-6):
+        super().__init__()
+        self.in_dim = in_dim
+        self.num_groups = num_groups if isinstance(num_groups, Dim) else Dim(num_groups, name="groups")
+        self.in_group_dim = in_dim.ceildiv_left(num_groups)
+        self.eps = eps
+        self.scale = rf.Parameter([self.in_dim] if isinstance(self.in_dim, rf.Dim) else self.in_dim)
+        self.scale.initial = 1.0
+        self.bias = rf.Parameter(self.scale.dims)
+        self.bias.initial = 0.0
+
+    def __call__(self, x: Tensor) -> Tensor:
+        x = rf.split_dims(x, axis=self.in_dim, dims=[self.num_groups, self.in_group_dim])
+        mean, variance = rf.moments(x, axis=self.in_group_dim)
+        norm_x = (x - mean) * rf.rsqrt(variance + self.eps)
+        norm_x, _ = rf.merge_dims(norm_x, dims=[self.num_groups, self.in_group_dim], out_dim=self.in_dim)
+        return norm_x * self.scale + self.bias
+
+
 class BatchNorm(rf.Module):
     """
     Batch normalization. https://arxiv.org/abs/1502.03167

{returnn-1.20240617.171744 → returnn-1.20240618.151147}/returnn/frontend/run_ctx.py

@@ -14,7 +14,7 @@ import returnn.frontend as rf
 from . import _backend
 
 
-__all__ = ["RunCtx", "Loss", "get_run_ctx", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
+__all__ = ["RunCtx", "Loss", "get_run_ctx", "get_run_ctx_step", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
 
 
 _run_ctx = None  # type: Optional[RunCtx]
@@ -134,6 +134,12 @@ class RunCtx:
         """
         return self._step
 
+    def get_step_tensor(self) -> Tensor:
+        """
+        :return: step as tensor
+        """
+        return rf.convert_to_tensor(self.step)
+
     @property
     def epoch(self) -> Union[int, Tensor]:
         """
@@ -141,6 +147,12 @@ class RunCtx:
         """
         return self._epoch
 
+    def get_epoch_tensor(self) -> Tensor:
+        """
+        :return: epoch as tensor
+        """
+        return rf.convert_to_tensor(self.epoch)
+
     def mark_as_loss(
         self,
         loss: Union[Tensor, Any],
@@ -350,6 +362,13 @@ class RunCtx:
         return loss
 
 
+def get_run_ctx_step() -> Tensor:
+    """
+    :return: shortcut for ``get_run_ctx().get_step_tensor()``
+    """
+    return get_run_ctx().get_step_tensor()
+
+
 @dataclass
 class Loss:
     """

returnn-1.20240618.151147/returnn/frontend/stepwise_scheduler.py

@@ -0,0 +1,31 @@
+"""
+Stepwise scheduler, e.g. for learning rate or other hyperparameters.
+
+All these modules will accept any args/kwargs but leave them unused,
+and instead uses :func:`get_run_ctx` to get the current train step from the current run context.
+"""
+
+from __future__ import annotations
+from returnn.tensor import Tensor
+import returnn.frontend as rf
+from .piecewise_linear import PiecewiseLinear
+
+
+__all__ = ["PiecewiseLinearStepwiseScheduler"]
+
+
+class PiecewiseLinearStepwiseScheduler(PiecewiseLinear):
+    """
+    Piecewise linear scheduler based on the current global train step.
+
+    Example::
+
+        scheduler = PiecewiseLinearStepwiseScheduler(
+            {0: 1.0, 10000: 0.1, 20000: 0.01}
+        )
+
+    This will start with 1.0, and then linearly decay to 0.1 at step 10000, and then to 0.01 at step 20000.
+    """
+
+    def __call__(self, *args, **kwargs) -> Tensor:
+        return super().__call__(rf.get_run_ctx_step())

{returnn-1.20240617.171744 → returnn-1.20240618.151147/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240617.171744 → returnn-1.20240618.151147}/returnn.egg-info/SOURCES.txt

@@ -162,6 +162,7 @@ returnn/frontend/_utils.py
 returnn/frontend/array_.py
 returnn/frontend/attention.py
 returnn/frontend/backend.py
+returnn/frontend/build_from_dict.py
 returnn/frontend/cond.py
 returnn/frontend/const.py
 returnn/frontend/container.py
@@ -191,6 +192,7 @@ returnn/frontend/reduce.py
 returnn/frontend/run_ctx.py
 returnn/frontend/signal.py
 returnn/frontend/state.py
+returnn/frontend/stepwise_scheduler.py
 returnn/frontend/tensor_array.py
 returnn/frontend/types.py
 returnn/frontend/_native/__init__.py

returnn-1.20240617.171744/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240617.171744'
-long_version = '1.20240617.171744+git.ef21456'