returnn 1.20240830.140746__tar.gz → 1.20240905.105440__tar.gz

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/PKG-INFO

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

returnn-1.20240905.105440/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240905.105440'
+long_version = '1.20240905.105440+git.d79c168'

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/datasets/distrib_files.py

@@ -317,15 +317,36 @@ class DistributeFilesDataset(CachedDataset2):
         return True
 
     def _get_sub_dataset_dict(self, files: List[FileTree]) -> Dict[str, Any]:
+        import tree
+
         dataset_dict = self.get_sub_epoch_dataset(files)
         dataset_dict = extend_dataset_dict_from_parent_dataset(dataset_dict, parent_dataset=self)
-        if dataset_dict.get("partition_epoch", 1) != 1:
-            raise ValueError(f"{self}: sub dataset should not have partition_epoch, got: {dataset_dict}")
-        if "seq_ordering" not in dataset_dict and "seq_order_control_dataset" not in dataset_dict:
+
+        flat_sub_dset = tree.flatten_with_path(dataset_dict)
+
+        part_epoch_cfg = next(
+            ((path, v) for path, v in flat_sub_dset if path[-1] == "partition_epoch" and v != 1), None
+        )
+        if part_epoch_cfg is not None:
+            path, subeps = part_epoch_cfg
             raise ValueError(
-                f"{self}: sub dataset should have explicit seq_ordering "
-                f"(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
+                f"{self}: sub dataset should not have partition_epoch, "
+                f'but got "partition_epoch": {subeps} at {".".join(path)} in {dataset_dict}.'
             )
+
+        # Heuristic check for well-definedness of seq ordering. Might need to be extended in the
+        # future if there are other ways of defining a seq order than the ones below.
+        if (
+            not any(path[-1] == "seq_ordering" for path, _ in flat_sub_dset)
+            and not any(path[-1] == "seq_order_control_dataset" for path, _ in flat_sub_dset)
+            and not any(path[-1] == "map_seq_stream" for path, _ in flat_sub_dset)
+        ):
+            raise ValueError(
+                f"{self}: there should be an explicit seq_ordering somewhere in the sub dataset "
+                f"(or seq_order_control_dataset for MetaDataset or map_seq_stream for PostprocessingDataset), "
+                f"but found none in {dataset_dict}."
+            )
+
         return dataset_dict
 
     @staticmethod

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/frontend/_numpy_backend.py

@@ -153,7 +153,10 @@ class NumpyBackend(Backend[numpy.ndarray]):
             op = NumpyBackend._CombineKindMap.get(kind)
             if not op:
                 raise ValueError(f"RF NumpyBackend: combine kind {kind!r} not supported")
-        return op(a, b)
+        res = op(a, b)
+        if not isinstance(res, numpy.ndarray):
+            res = numpy.array(res)
+        return res
 
     @staticmethod
     def range_over_dim(dim: Dim, *, dtype: Optional[str] = None, device: Optional[str] = None) -> Tensor[numpy.ndarray]:
@@ -211,3 +214,14 @@ class NumpyBackend(Backend[numpy.ndarray]):
             sparse_dim=source.sparse_dim,
         )
         return res
+
+    @staticmethod
+    def activation_raw(raw_tensor: numpy.ndarray, func: str) -> numpy.ndarray:
+        """
+        :param raw_tensor:
+        :param func: "tanh", "sigmoid", "relu", ...
+        :return: raw tensor with elementwise activation applied
+        """
+        if func == "relu":
+            return numpy.array(numpy.maximum(raw_tensor, 0))
+        raise NotImplementedError("NumpyBackend: activation %r not implemented" % func)

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/tensor/_dim_extra.py

@@ -4,8 +4,9 @@ or just rarely used attribs, such that we can save memory for the common case.
 """
 
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional, Union, Any, Tuple, Sequence, Dict, List, Set, Callable
+from typing import TYPE_CHECKING, Optional, Union, Any, Tuple, Sequence, MutableMapping, Dict, List, Set, Callable
 import operator
+import weakref
 
 from returnn.util.basic import Entity
 from returnn.util import basic as util
@@ -118,7 +119,7 @@ class _DimExtra:
         self.same_for_batch_ctx = {}  # type: Dict[Tuple[BatchInfo,Optional[ControlFlowContext]],_d.Dim]
         self.cache_dyn_size_ext_dev = {}  # type: Dict[str,_t.Tensor]  # device -> dyn_size_ext
         self.cache_seq_mask: Dict[Tuple[str, Optional[Tuple[Dim, ...]]], _t.Tensor] = {}  # (dev,dim_order) -> seq_mask
-        self.cache_dim_math: Dict[Tuple[str, Union[Dim, int]], Dim] = {}  # op (add,sub,...), operand -> Dim
+        self.cache_dim_math = _CacheDimMath()  # op (add,sub,...), operand -> Dim
 
     def __getstate__(self):
         d = vars(self).copy()
@@ -389,6 +390,10 @@ class _DimMixin:
             if dim_extra:
                 # Any dims via dim math could also contain raw tensors,
                 # so iterate through them.
+                if dim.dyn_size_ext is not None or dim.dimension is None:
+                    dim_extra.cache_dim_math.clear()
+                else:
+                    dim_extra.cache_dim_math.clear_dynamic()
                 queue += dim_extra.cache_dim_math.values()
                 if dim_extra.same_as:
                     queue.append(dim_extra.same_as)
@@ -2873,6 +2878,123 @@ def dim_cmp_value(obj):
     return obj
 
 
+class _CacheDimMath:
+    """op (add,sub,...), operand -> Dim"""
+
+    class _OperandCache:
+        def __init__(self):
+            self.dims: MutableMapping[Dim, Dim] = weakref.WeakKeyDictionary()
+            self.statics: Dict[int, Dim] = {}
+
+    def __init__(self):
+        self._ops: Dict[str, _CacheDimMath._OperandCache] = {}
+
+    def __repr__(self):
+        return "_CacheDimMath({%s})" % ", ".join("%r: %r" % (k, v) for k, v in self.items())
+
+    def _get_op_dict(self, __key: Tuple[str, Union[Dim, int]]) -> _OperandCache:
+        if __key[0] in self._ops:
+            return self._ops[__key[0]]
+        else:
+            op_dict = self._OperandCache()
+            self._ops[__key[0]] = op_dict
+            return op_dict
+
+    def __setitem__(self, __key: Tuple[str, Union[Dim, int]], __value: Dim):
+        op_dict = self._get_op_dict(__key)
+        if isinstance(__key[1], int):
+            value_dict = op_dict.statics
+        else:
+            value_dict = op_dict.dims
+        if __key[1] in value_dict:
+            value_dict[__key[1]] = __value
+            return
+        if len(value_dict) >= 5:
+            # Just to avoid memory leaks.
+            value_dict.clear()
+        value_dict[__key[1]] = __value
+
+    def __delitem__(self, __key: Tuple[str, Union[Dim, int]]):
+        op_dict = self._ops[__key[0]]
+        if isinstance(__key[1], int):
+            del op_dict.statics[__key[1]]
+        else:
+            del op_dict.dims[__key[1]]
+
+    def __getitem__(self, __key: Tuple[str, Union[Dim, int]]) -> Dim:
+        op_dict = self._ops[__key[0]]
+        if isinstance(__key[1], int):
+            return op_dict.statics[__key[1]]
+        else:
+            return op_dict.dims[__key[1]]
+
+    def __contains__(self, __key: Tuple[str, Union[Dim, int]]) -> bool:
+        op_dict = self._ops.get(__key[0])
+        if not op_dict:
+            return False
+        if isinstance(__key[1], int):
+            return __key[1] in op_dict.statics
+        else:
+            return __key[1] in op_dict.dims
+
+    def get(self, __key: Tuple[str, Union[Dim, int]], default: Optional[Dim] = None) -> Optional[Dim]:
+        """get"""
+        op_dict = self._ops.get(__key[0])
+        if not op_dict:
+            return default
+        if isinstance(__key[1], int):
+            return op_dict.statics.get(__key[1], default)
+        else:
+            return op_dict.dims.get(__key[1], default)
+
+    def setdefault(self, __key: Tuple[str, Union[Dim, int]], __value: Dim):
+        """setdefault"""
+        existing = self.get(__key)
+        if existing is not None:
+            return existing
+        self[__key] = __value
+        return __value
+
+    def clear(self):
+        """clear"""
+        self._ops.clear()
+
+    def clear_dynamic(self):
+        """clear dynamic part"""
+        for op_dict in self._ops.values():
+            for k, v in list(op_dict.dims.items()):
+                if v.dyn_size_ext is not None or v.dimension is None:
+                    del op_dict.dims[k]
+
+    def __len__(self):
+        count = 0
+        for op_dict in self._ops.values():
+            count += len(op_dict.statics)
+            count += len(op_dict.dims)
+        return count
+
+    def items(self):
+        """items"""
+        for op_name, op_dict in self._ops.items():
+            for key, value in op_dict.statics.items():
+                yield (op_name, key), value
+            for key, value in op_dict.dims.items():
+                yield (op_name, key), value
+
+    def keys(self):
+        """keys"""
+        for k, v in self.items():
+            yield k
+
+    def values(self):
+        """values"""
+        for k, v in self.items():
+            yield v
+
+    def __iter__(self):
+        yield from self.keys()
+
+
 def _behavior_version_reset_callback():
     # Reset things we did in _handle_new_min_version.
     _DimMixin._SimpleEquality = False

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/tf/frontend_layers/config_entry_points.py

@@ -118,6 +118,9 @@ def get_net_dict(
             # but now the TF engine actually wants to have Tensor[tf.Tensor].
             # Reset it now. The TF engine should redefine it again.
             elem.reset_batch_and_raw()
+        elif isinstance(elem, set):
+            # map_structure does not recurse into sets.
+            nest.map_structure(_cleanup_net_dict_value, sorted(list(elem)))
         return elem
 
     # Do some cleanup.

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/tf/util/data.py

@@ -339,7 +339,7 @@ class BatchInfo:
 
         # Ok, need to extend.
         global_batch_dims = [dim for dim in all_virtual_dims if isinstance(dim, BatchInfo.GlobalBatchDim)]
-        assert len(global_batch_dims) == 1
+        assert len(global_batch_dims) == 1, f"got global_batch_dims={global_batch_dims!r}"
         global_batch_dim = global_batch_dims[0]
         assert base.virtual_dims == [global_batch_dim]
         beams = [dim for dim in all_virtual_dims if isinstance(dim, BatchInfo.BeamDim)]

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/torch/engine.py

@@ -3,13 +3,15 @@ Main engine for PyTorch
 """
 
 from __future__ import annotations
-from typing import Optional, Any, Union, Callable, Dict
+from typing import Optional, Any, Union, Callable, Dict, Set
 from contextlib import nullcontext
 
 import gc
 import os
 import time
 import socket
+import fnmatch
+import re
 
 import torch
 import torch.distributed
@@ -41,6 +43,7 @@ from .data import extern_data as extern_data_util
 from .data.queued_data_iter import QueuedDataIter
 from .frontend.bridge import rf_module_to_pt_module
 from .util import diagnose_gpu
+from .util import module as util_module
 from .distributed import DistributedContext, get_ctx as dist_get_ctx
 
 
@@ -83,6 +86,7 @@ class Engine(EngineBase):
             self._forward_step_expected_outputs = TensorDict()
             self._forward_step_expected_outputs.update(self.config.typed_value("model_outputs"), auto_convert=True)
         self._save_model_epoch_interval = 1
+        self._ignore_param_set: Set[str] = set()  # for the updater and for saving the model checkpoint
         self._updater = None  # type: Optional[Updater]
 
         self._use_autocast = False
@@ -721,6 +725,7 @@ class Engine(EngineBase):
 
         self._create_model(epoch=epoch, step=step)
 
+        self._ignore_param_set.clear()
         loaded_state_keys = set()
         missing_keys = set()
         unexpected_keys = set()
@@ -745,6 +750,7 @@ class Engine(EngineBase):
 
         preload_from_files = self.config.typed_value("preload_from_files", {})
         if preload_from_files:
+            model_state_keys_set = set(self._pt_model.state_dict().keys())
             # see `preload_from_files` in tf engine and `returnn.tf.network.CustomCheckpointLoader`
             # We use the reversed sorted order here to achieve consistent behavior with the TF engine.
             # There, the keys are used in sorted order but if a variable is loaded,
@@ -754,12 +760,39 @@ class Engine(EngineBase):
             # In order to get consistent behavior, we use the reversed order.
             for preload_key, opts in reversed(sorted(preload_from_files.items())):
                 assert isinstance(opts, dict) and "filename" in opts
-                if opts.get("init_for_train", False):
-                    if not is_first_train_epoch:
-                        continue
+                init_for_train = opts.get("init_for_train", False)
+                if init_for_train:
+                    if isinstance(init_for_train, str) and init_for_train == "always":
+                        # No matter if this is the first train epoch
+                        # or training with loading some prev epoch,
+                        # those parameters will always be loaded via preload_from_files,
+                        # and thus also not stored in our own checkpoint.
+                        pass
+                    elif isinstance(init_for_train, bool) and init_for_train:
+                        if not is_first_train_epoch:
+                            continue
+                    else:
+                        raise ValueError(
+                            f"preload key {preload_key}:"
+                            f" invalid init_for_train value {init_for_train!r} (type {type(init_for_train).__name__})"
+                        )
                 else:  # default: init for recog
                     if is_training:
                         continue
+                if opts["filename"] is None:
+                    print(f"Pre-load (initialize) weights for key '{preload_key}'", file=log.v3)
+                    pattern = opts["pattern"]
+                    match = re.compile(fnmatch.translate(pattern)).match
+                    remove = []
+                    for name in self._pt_model.state_dict().keys():
+                        if match(name) and name in missing_keys:
+                            remove.append(name)
+                    if remove:
+                        print(f"Randomly initialize params: {remove}", file=log.v3)
+                        missing_keys.difference_update(remove)
+                    else:
+                        print("(No relevant parameters matching.)", file=log.v3)
+                    continue
                 print(f"Pre-load weights for key '{preload_key}' from {opts['filename']}", file=log.v3)
                 preload_model_state = torch.load(opts["filename"])
                 if opts.get("checkpoint_key", "model") is not None:
@@ -786,6 +819,8 @@ class Engine(EngineBase):
                         preload_model_state.pop(key)
                 for new_name, name_in_checkpoint in opts.get("var_name_mapping", {}).items():
                     preload_model_state[new_name] = preload_model_state.pop(name_in_checkpoint)
+                if init_for_train == "always":
+                    self._ignore_param_set.update(set(preload_model_state.keys()).intersection(model_state_keys_set))
                 missing_keys_preload, unexpected_keys_preload = self._pt_model.load_state_dict(
                     preload_model_state, strict=False
                 )
@@ -797,7 +832,7 @@ class Engine(EngineBase):
                 if opts.get("prefix", ""):
                     prefix_keys = [key for key in self._pt_model.state_dict() if key.startswith(opts.get("prefix", ""))]
                 else:
-                    prefix_keys = self._pt_model.state_dict().keys()
+                    prefix_keys = model_state_keys_set
                 missing_keys_preload = (
                     set(prefix_keys).intersection(set(missing_keys_preload)).difference(loaded_state_keys)
                 )
@@ -816,6 +851,9 @@ class Engine(EngineBase):
                     )
                     unexpected_keys.update(unexpected_keys_preload)
 
+        if self._ignore_param_set:
+            util_module.convert_parameters_to_buffers(self._pt_model, self._ignore_param_set, persistent=False)
+
         if missing_keys:
             raise Exception(
                 "\n".join(
@@ -913,9 +951,16 @@ class Engine(EngineBase):
         tmp_filename = filename + ".tmp_write"
         if os.path.exists(tmp_filename):
             os.unlink(tmp_filename)
+        state_dict = self._pt_model.state_dict()
+        if self._ignore_param_set:
+            # Do some extra check that we don't save the ignored parameters.
+            # Should not be in the state_dict anymore because we should have converted them to buffers
+            # via util_module.convert_parameters_to_buffers before.
+            remaining = set(state_dict.keys()).intersection(self._ignore_param_set)
+            assert not remaining, f"_save_model: found remaining params in state_dict to ignore: {remaining}"
         torch.save(
             {
-                "model": self._pt_model.state_dict(),
+                "model": state_dict,
                 "epoch": self.epoch,
                 "step": self.global_train_step,
                 "effective_learning_rate": self._updater.get_effective_learning_rate() if self._updater else None,

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn/torch/frontend/bridge.py

@@ -149,13 +149,16 @@ class RFModuleAsPTModule(torch.nn.Module):
         for name, rf_param in self._rf_module.named_parameters(recurse=False):
             pt_param = getattr(self, name)
             if rf_param.auxiliary and self._aux_params_as_buffers:
-                assert isinstance(pt_param, torch.Tensor)  # but not torch.nn.Parameter
-                # See similar logic in torch.nn.Module._apply.
-                pt_param = torch.nn.Parameter(pt_param, pt_param.requires_grad)
-            else:
-                assert isinstance(pt_param, torch.nn.Parameter), (
-                    f"{self}.{name} is not a Parameter" f" but {type(pt_param).__name__}"
-                )
+                if not isinstance(pt_param, torch.nn.Parameter):
+                    assert isinstance(pt_param, torch.Tensor)  # but not torch.nn.Parameter
+                    # See similar logic in torch.nn.Module._apply.
+                    pt_param = torch.nn.Parameter(pt_param, pt_param.requires_grad)
+            # Otherwise, we do not care whether it is a torch.nn.Parameter or not.
+            # Its type might have changed due to convert_parameters_to_buffers.
+            # Just make sure it is a tensor.
+            assert isinstance(pt_param, torch.Tensor)
+            # noinspection PyProtectedMember
+            rf_param.dtype = rf_param._raw_backend.get_dtype_name_raw(pt_param)  # dtype might have changed
             rf_param.raw_tensor = pt_param
 
     def register_parameter(self, name: str, param: Optional[torch.nn.Parameter]) -> None:

returnn-1.20240905.105440/returnn/torch/util/module.py

@@ -0,0 +1,43 @@
+"""
+Utils for modules
+"""
+
+from __future__ import annotations
+from typing import Collection
+import torch
+
+
+def convert_parameters_to_buffers(
+    module: torch.nn.Module, parameter_names: Collection[str], *, deep: bool = True, persistent: bool
+):
+    """
+    :param module:
+    :param parameter_names:
+    :param deep: parameter_name can contain '.' to access submodules
+    :param persistent: whether the buffer is persistent. if True, the buffer will be saved to the state_dict.
+        passed to module.register_buffer.
+    """
+    for parameter_name in parameter_names:
+        convert_parameter_to_buffer(module, parameter_name, deep=deep, persistent=persistent)
+
+
+def convert_parameter_to_buffer(module: torch.nn.Module, parameter_name: str, *, deep: bool = True, persistent: bool):
+    """
+    :param module:
+    :param parameter_name:
+    :param deep: parameter_name can contain '.' to access submodules
+    :param persistent: whether the buffer is persistent. if True, the buffer will be saved to the state_dict.
+        passed to module.register_buffer.
+    """
+    if "." in parameter_name:
+        if not deep:
+            raise ValueError("parameter_name can't contain '.' when deep is False")
+        module_path, _, parameter_name = parameter_name.rpartition(".")
+        module = module.get_submodule(module_path)
+
+    parameter = getattr(module, parameter_name)
+    if not isinstance(parameter, torch.nn.Parameter):
+        raise ValueError(f"{parameter_name} is not a torch.nn.Parameter, got type {type(parameter).__name__}")
+    delattr(module, parameter_name)
+    parameter.requires_grad = False
+    module.register_buffer(parameter_name, parameter, persistent=persistent)

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/returnn.egg-info/SOURCES.txt

@@ -302,6 +302,7 @@ returnn/torch/util/__init__.py
 returnn/torch/util/array_.py
 returnn/torch/util/diagnose_gpu.py
 returnn/torch/util/gradient_checkpoint.py
+returnn/torch/util/module.py
 returnn/torch/util/scaled_gradient.py
 returnn/util/__init__.py
 returnn/util/basic.py

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/tests/rf_utils.py

@@ -185,6 +185,9 @@ def _run_model_net_dict_tf(
 
     from returnn.tf.frontend_layers.config_entry_points import get_net_dict
 
+    # noinspection PyProtectedMember
+    from returnn.frontend import _backend
+
     config = Config(
         {
             "debug_runtime_sanity_checks": True,
@@ -203,6 +206,7 @@ def _run_model_net_dict_tf(
         outputs_layers = rf.get_run_ctx().outputs
         print("*** outputs:", outputs_layers)
 
+        _backend.select_backend_tf()
         net = TFNetwork(config=config, train_flag=False)
         net.construct_from_dict(net_dict)
 

{returnn-1.20240830.140746 → returnn-1.20240905.105440}/tests/test_rf_array.py

@@ -238,7 +238,7 @@ def test_pad_time_right():
     assert data_.dims == (batch_dim, time_dim, in_dim)
     new_time_dim = out_.dims[1]
     assert out_.dims == (batch_dim, new_time_dim, in_dim) and new_time_dim != time_dim
-    assert new_time_dim == time_dim + 1  # math dim... not really necessary check here...
+    # assert new_time_dim == time_dim + 1  # math dim... not really necessary check here...
     assert time_dim.dyn_size_ext.dims == new_time_dim.dyn_size_ext.dims == (batch_dim,)
     batch_size = batch_dim.get_dim_value()
     assert batch_size > 1