returnn 1.20240626.175748__tar.gz → 1.20240628.104058__tar.gz

{returnn-1.20240626.175748 → returnn-1.20240628.104058}/PKG-INFO

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

returnn-1.20240628.104058/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240628.104058'
+long_version = '1.20240628.104058+git.32c390d'

{returnn-1.20240626.175748 → returnn-1.20240628.104058}/returnn/datasets/distrib_files.py

@@ -138,6 +138,7 @@ class DistributeFilesDataset(CachedDataset2):
         buffer_size: int = 1,
         distrib_shard_files: bool = False,
         _meta_info_cache: Optional[Dict[str, Any]] = None,
+        _distrib_info: Optional[Dict[str, int]] = None,
         **kwargs,
     ):
         """
@@ -149,6 +150,7 @@ class DistributeFilesDataset(CachedDataset2):
         :param distrib_shard_files: set to true to shard the data across worker processes in
             distributed training scenaria
         :param _meta_info_cache: for internal use
+        :param _distrib_info: for internal use
         """
         super().__init__(**kwargs)
         self.files = files
@@ -159,12 +161,24 @@ class DistributeFilesDataset(CachedDataset2):
         self._file_sizes: Optional[Dict[str, int]] = None  # key -> size. for equal distribution across sub epochs
         self._data_keys: Optional[List[str]] = None
         self._num_seqs: Optional[int] = None
-        self.distrib_shard_files = distrib_shard_files
-        self._shard_index, self._num_shards = _get_rank_and_size() if distrib_shard_files else (0, 1)
 
         self._workers: Dict[int, _WorkerProcParent] = {}  # epoch -> worker
         self._files_order_cache: Dict[int, List[List[FileTree]]] = {}  # full epoch (0-indexed) -> files order
 
+        self.distrib_shard_files = distrib_shard_files
+        if distrib_shard_files:
+            if _distrib_info:
+                # If we're in a child process `_get_rank_and_size()` no longer works,
+                # so we pass the info about the shards via a pickled property.
+                # See also Dataset.__reduce__.
+                self._shard_index = _distrib_info["shard_index"]
+                self._num_shards = _distrib_info["num_shards"]
+            else:
+                self._shard_index, self._num_shards = _get_rank_and_size()
+        else:
+            self._shard_index = 0
+            self._num_shards = 1
+
         if _meta_info_cache:
             # This allows to skip the lazy init in self.initialize().
             # This can be used when pickling/deepcopying an instance of this dataset,
@@ -183,6 +197,10 @@ class DistributeFilesDataset(CachedDataset2):
         self._lazy_init_num_outputs()
         super().initialize()
 
+    @property
+    def _distrib_info(self):
+        return {"num_shards": self._num_shards, "shard_index": self._shard_index}
+
     @property
     def _meta_info_cache(self):
         if not self.num_outputs:
@@ -220,9 +238,10 @@ class DistributeFilesDataset(CachedDataset2):
         }
 
     def __del__(self):
-        for k, worker in self._workers.items():
-            try_run(worker.exit, kwargs={"join": False})
-        self._workers.clear()
+        if hasattr(self, "_workers"):  # might not be set after an early exception in __init__
+            for k, worker in self._workers.items():
+                try_run(worker.exit, kwargs={"join": False})
+            self._workers.clear()
 
     def init_seq_order(self, epoch: Optional[int] = None, seq_list=None, seq_order=None) -> bool:
         """

{returnn-1.20240626.175748 → returnn-1.20240628.104058}/returnn/frontend/build_from_dict.py

@@ -3,8 +3,10 @@ Construct modules (or other objects) from dictionaries.
 """
 
 from __future__ import annotations
-from typing import Union, Any, Type, Dict
+from types import FunctionType, BuiltinFunctionType
+from typing import Union, Any, Type, Callable, Dict
 import importlib
+import functools
 import returnn.frontend as rf
 
 
@@ -20,17 +22,23 @@ def build_from_dict(d: Dict[str, Any], *args, **kwargs) -> Union[rf.Module, Any]
     :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)
+    :return: ``cls(*args, **d, **kwargs)`` (after "class" was popped from ``d``)
     """
     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)
+    if isinstance(cls, (FunctionType, BuiltinFunctionType)):
+        if args or kwargs or d:
+            return functools.partial(cls, *args, **d, **kwargs)
+        return cls
+    if not isinstance(cls, type):
+        raise ValueError(f"build_from_dict: Expected class, got {cls!r} for class_name {cls_name!r}")
     return cls(*args, **d, **kwargs)
 
 
-def build_dict(cls: Type, **kwargs) -> Dict[str, Any]:
+def build_dict(cls: Union[Type, FunctionType, BuiltinFunctionType, Callable[..., Any]], **kwargs) -> Dict[str, Any]:
     """
     Build a dictionary for :func:`build_from_dict`.
     The class name is stored in the `"class"` key.
@@ -40,7 +48,15 @@ def build_dict(cls: Type, **kwargs) -> Dict[str, Any]:
     (e.g. for Sisyphus :func:`sis_hash_helper`)
     which should not change if the class is renamed or moved
     to keep the hash stable.
+
+    :param cls: class or function. actually we do not expect any other Callable here.
+        It's just Callable in the type hint because MyPy/PyCharm don't correctly support FunctionType
+            (https://github.com/python/mypy/issues/3171).
+    :param kwargs: other kwargs put into the dict. expect to contain only other serializable values.
+    :return: build dict. can easily be serialized. to be used with :func:`build_from_dict`.
     """
+    if not isinstance(cls, (type, FunctionType, BuiltinFunctionType)):
+        raise TypeError(f"build_dict: Expected class or function, got: {cls!r}")
     return {"class": _get_cls_name(cls), **kwargs}
 
 

{returnn-1.20240626.175748 → returnn-1.20240628.104058}/returnn/frontend/encoder/conformer.py

@@ -21,7 +21,14 @@ class ConformerPositionwiseFeedForward(rf.Module):
         FF -> Activation -> Dropout -> FF
     """
 
-    def __init__(self, out_dim: Dim, *, ff_dim: Dim, dropout: float, activation: Callable[[Tensor], Tensor]):
+    def __init__(
+        self,
+        out_dim: Dim,
+        *,
+        ff_dim: Dim,
+        dropout: float,
+        activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module],
+    ):
         """
         :param out_dim: output feature dimension
         :param ff_dim: dimension of the feed-forward layers
@@ -33,7 +40,9 @@ class ConformerPositionwiseFeedForward(rf.Module):
         self.out_dim = out_dim
         self.dropout = dropout
         self.dropout_broadcast = rf.dropout_broadcast_default()
-        if not callable(activation):
+        if isinstance(activation, dict):
+            activation = rf.build_from_dict(activation)
+        elif not callable(activation):
             raise TypeError(f"{self}: unexpected activation type {activation!r}")
         self.activation = activation
 
@@ -180,7 +189,7 @@ class ConformerEncoderLayer(rf.Module):
         out_dim: Dim = Dim(512, name="conformer-enc-default-out-dim"),
         *,
         ff_dim: Dim = NotSpecified,
-        ff_activation: Callable[[Tensor], Tensor] = rf.swish,
+        ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = rf.swish,
         dropout: float = 0.1,
         conv_kernel_size: int = 32,
         conv_norm: Union[rf.BatchNorm, type, Dict[str, Any], Any] = NotSpecified,
@@ -307,7 +316,7 @@ class ConformerEncoder(ISeqDownsamplingEncoder):
         input_layer: Union[ConformerConvSubsample, ISeqDownsamplingEncoder, rf.Module, Any],
         input_dropout: float = 0.1,
         ff_dim: Dim = NotSpecified,
-        ff_activation: Callable[[Tensor], Tensor] = rf.swish,
+        ff_activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module] = rf.swish,
         dropout: float = 0.1,
         conv_kernel_size: int = 32,
         conv_norm: Union[rf.BatchNorm, type, Dict[str, Any], Any] = NotSpecified,

{returnn-1.20240626.175748 → returnn-1.20240628.104058}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240626.175748 → returnn-1.20240628.104058}/tests/test_rf_base.py

@@ -397,3 +397,43 @@ def test_rf_range_over_dim():
         rf.range_over_dim(time_dim, dtype="float32").mark_as_output("range_float", shape=[time_dim])
 
     run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+def test_build_dict():
+    d = rf.build_dict(rf.BatchNorm, use_mask=True)
+    assert d == {"class": "rf.BatchNorm", "use_mask": True}
+
+
+def test_build_dict_func():
+    d = rf.build_dict(rf.relu)
+    assert d == {"class": "rf.relu"}
+
+
+def test_build_from_dict():
+    rf.select_backend_torch()
+    in_dim = Dim(7, name="in")
+    mod = rf.build_from_dict({"class": "rf.BatchNorm", "use_mask": True}, in_dim, eps=1e-4)
+    assert isinstance(mod, rf.BatchNorm)
+    assert mod.in_dim == in_dim
+    assert mod.use_mask
+    assert mod.eps == 1e-4
+
+
+def test_build_from_dict_func():
+    import functools
+
+    func = rf.build_from_dict({"class": "rf.relu"})
+    assert func is rf.relu
+    func = rf.build_from_dict({"class": "rf.combine", "kind": "+", "b": 1})
+    assert isinstance(func, functools.partial)
+    assert func.func is rf.combine
+    assert func.keywords == {"kind": "+", "b": 1}
+
+
+def test_build_from_dict_func_native():
+    from types import BuiltinFunctionType
+
+    rf.select_backend_torch()  # enables some of the native optimizations
+    assert isinstance(rf.combine, BuiltinFunctionType)  # due to native optimizations
+    func = rf.build_from_dict({"class": "rf.combine"})
+    assert func is rf.combine

returnn-1.20240626.175748/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240626.175748'
-long_version = '1.20240626.175748+git.80f52ec'