returnn 1.20230403.124714__tar.gz → 1.20230403.211148__tar.gz

{returnn-1.20230403.124714/returnn.egg-info → returnn-1.20230403.211148}/PKG-INFO

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

returnn-1.20230403.211148/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230403.211148'
+long_version = '1.20230403.211148+git.aa0221f'

{returnn-1.20230403.124714 → returnn-1.20230403.211148}/returnn/frontend/run_ctx.py

@@ -9,11 +9,11 @@ or forwarding loop.
 from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, Dict
 from dataclasses import dataclass
-from returnn.tensor import Tensor, Dim
+from returnn.tensor import Tensor, Dim, TensorDict
 import returnn.frontend as rf
 
 
-__all__ = ["RunCtx", "Loss", "Output", "get_run_ctx", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
+__all__ = ["RunCtx", "Loss", "get_run_ctx", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
 
 
 _run_ctx = None  # type: Optional[RunCtx]
@@ -75,7 +75,7 @@ class RunCtx:
         """
         self.stage = stage
         self.losses = {}  # type: Dict[str, Loss]
-        self.outputs = {}  # type: Dict[str, Output]
+        self.outputs = TensorDict()
 
     def mark_as_loss(
         self,
@@ -141,7 +141,7 @@ class RunCtx:
             custom_inv_norm_factor=custom_inv_norm_factor,
         )
 
-    def mark_as_output(self, tensor: Union[Tensor, Any], name: str, *, shape: Optional[Sequence[int]] = None) -> None:
+    def mark_as_output(self, tensor: Union[Tensor, Any], name: str, *, dims: Optional[Sequence[int]] = None) -> None:
         """
         Mark this as an output.
         This has the effect that RETURNN will in any case construct the corresponding layer.
@@ -153,7 +153,7 @@ class RunCtx:
 
         :param tensor:
         :param name:
-        :param shape: this specifies the order of the dims of the output, such that it is well-defined
+        :param dims: this specifies the order of the dims of the output, such that it is well-defined
             for some external application.
             If not specified, we try to infer BTF or BF as default, if that works, otherwise it will be an error.
         """
@@ -161,7 +161,32 @@ class RunCtx:
         if not isinstance(tensor, Tensor):
             tensor = rf.convert_to_tensor(tensor)
         assert name not in self.outputs
-        self.outputs[name] = Output(tensor=tensor, name=name, shape=shape)
+        if dims is None:
+            rem_dims = list(tensor.dims)
+            dims = []
+            if tensor.have_batch_axis():
+                rem_dims.remove(tensor.get_batch_dim_tag())
+                dims.append(tensor.get_batch_dim_tag())
+            if tensor.have_time_axis():
+                rem_dims.remove(tensor.get_time_dim_tag())
+                dims.append(tensor.get_time_dim_tag())
+            static_dims = [d for d in dims if d.is_static()]
+            if len(static_dims) > 1:
+                raise Exception(
+                    f"Cannot infer order of dims automatically for output {name!r}. Please specify a shape explicitly."
+                )
+            elif len(static_dims) == 1:
+                rem_dims.remove(static_dims[0])
+                dims.insert(0, static_dims[0])
+            if len(rem_dims) > 1:
+                raise Exception(
+                    f"Cannot infer order of dims automatically for output {name!r}. Please specify a shape explicitly."
+                )
+            elif len(rem_dims) == 1:
+                dims.append(rem_dims[0])
+        tensor = tensor.copy_transpose(dims, allow_int=False)
+        tensor = tensor.copy(name=name)
+        self.outputs.data[name] = tensor
 
     def mark_as_default_output(self, tensor: Union[Tensor, Any], *, shape: Optional[Sequence[Dim]] = None) -> None:
         """
@@ -173,7 +198,7 @@ class RunCtx:
         :param tensor:
         :param shape:
         """
-        self.mark_as_output(tensor, "output", shape=shape)
+        self.mark_as_output(tensor, "output", dims=shape)
 
     def total_loss(self) -> Union[Tensor, float]:
         """
@@ -233,16 +258,3 @@ class Loss:
         else:
             loss = self.get_summed_loss()
         return loss * self.scale
-
-
-@dataclass
-class Output:
-    """
-    Output via :func:`RunCtx.mark_as_output`.
-
-    We collect all relevant information here.
-    """
-
-    tensor: Tensor
-    name: str
-    shape: Optional[Sequence[Dim]] = None

{returnn-1.20230403.124714 → returnn-1.20230403.211148}/returnn/tensor/_tensor_extra.py

@@ -651,14 +651,16 @@ class _TensorMixin(_TensorMixinBase):
         assert self.time_dim_axis is not None
         return self.copy_move_axis(self.time_dim_axis, time_dim_axis)
 
-    def copy_transpose(self, perm) -> _t.Tensor:
+    def copy_transpose(self, perm: Sequence[Union[int, Dim]], *, allow_int: bool = True) -> _t.Tensor:
         """
-        :param list[int] perm: permutation of the axes, counted with batch-dim.
-          Maps the new axes to the old axes
+        :param perm: permutation of the axes. Maps the new axes to the old axes
+        :param allow_int: allow int as axis, otherwise only :class:`Dim`
         :return: copy of myself with permuted axes
         """
-        assert len(perm) == self.batch_ndim
-        assert set(perm) == set(range(self.batch_ndim))
+        assert len(perm) == self.batch_ndim, f"{self}: invalid perm {perm!r} length"
+        perm_ = perm
+        perm = [self.get_axis_from_description(a, allow_int=allow_int) for a in perm]
+        assert set(perm) == set(range(self.batch_ndim)), f"{self}: invalid perm {perm_!r} (axes: {perm!r})"
         if all(perm[axis] == axis for axis in range(self.batch_ndim)):
             return self.copy()
 
@@ -3156,7 +3158,7 @@ class _TensorMixin(_TensorMixinBase):
         """
         import returnn.frontend as rf
 
-        rf.get_run_ctx().mark_as_output(self, name=name, shape=shape)
+        rf.get_run_ctx().mark_as_output(self, name=name, dims=shape)
 
     def mark_as_default_output(self: Tensor, *, shape: Optional[Sequence[Dim]] = None) -> None:
         """

{returnn-1.20230403.124714 → returnn-1.20230403.211148}/returnn/tf/frontend_layers/config_entry_points.py

@@ -123,12 +123,12 @@ def get_net_dict(
         # Note that this logic might change.
         root_scope.marked_losses.append(loss_t)
 
-    for out in rf.get_run_ctx().outputs.values():
+    for out in rf.get_run_ctx().outputs.data.values():
         if out.name == "output" and out.name not in root_scope.children:
             layer = root_scope.get_child(out.name)
         else:
             layer = root_scope.get_new_child(suggested_name=out.name)
-        out_t = _utils.copy(out.tensor, name=layer)
+        out_t = _utils.copy(out, name=layer)
         if layer.name != "output":
             out_t.raw_tensor.layer_dict["is_output_layer"] = True
         root_scope.marked_outputs.append(out_t)

{returnn-1.20230403.124714 → returnn-1.20230403.211148}/returnn/torch/data/pipeline.py

@@ -18,17 +18,19 @@ However, having this separate pure PyTorch implementation is useful to allow to
 other PyTorch datasets more directly, including also HuggingFace datasets.
 """
 
+from __future__ import annotations
+from typing import List, Dict
 import sys
 from copy import deepcopy
 
-import numpy as np
+import numpy
 import torch
 import torch.utils.data
 
 from returnn.util.basic import NumbersDict
 
 
-def create_tensor(array: np.ndarray) -> torch.Tensor:
+def create_tensor(array: numpy.ndarray) -> torch.Tensor:
     """
     Adjust non-supported dtypes
 
@@ -36,14 +38,14 @@ def create_tensor(array: np.ndarray) -> torch.Tensor:
     """
     # The only supported PyTorch dtypes are:
     # float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.
-    if array.dtype == np.uint32:
-        array = np.asarray(array, dtype=np.int64)
+    if array.dtype == numpy.uint32:
+        array = numpy.asarray(array, dtype=numpy.int64)
     return torch.tensor(array)
 
 
-def collate_batch(batch):
+def collate_batch(batch: List[Dict[str, numpy.ndarray]]) -> Dict[str, torch.Tensor]:
     """
-    :param list[dict[str, numpy.ndarray]] batch:
+    :param batch:
     """
     assert isinstance(batch, list)
     assert batch, "batch is empty?"

{returnn-1.20230403.124714 → returnn-1.20230403.211148}/returnn/torch/engine.py

@@ -294,6 +294,59 @@ class Engine(EngineBase):
             print("Load model %s" % (filename,), file=log.v4)
             model_state = torch.load(filename)
             self._model.load_state_dict(model_state)
+        preload_from_files = self.config.typed_value("preload_from_files", {})
+        if preload_from_files:
+            # see `preload_from_files` in tf engine and `returnn.tf.network.CustomCheckpointLoader`
+            is_training = self.config.value("task", "train") == "train"
+            is_first_train_epoch = epoch == 1 and (
+                is_training or self.config.value("task", "train") == "initialize_model"
+            )
+            # 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,
+            # it will not be considered anymore afterwards.
+            # So the first occurrence is used.
+            # Here, we overwrite variables even if they have been loaded before.
+            # 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
+                else:  # default: init for recog
+                    if is_training:
+                        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", None) is not None:
+                    # This can be used if an external checkpoint saves a checkpoint a different structure that just the
+                    # model state dict. E.g., if a checkpoint is created using
+                    # `torch.save({"model": model.state_dict(), "optimizer": optimizer.state)_dict(), ...})`
+                    # we can set checkpoint_key = "model" to load the model.
+                    # Currently, this only supports single level dicts, but it could be extended if needed.
+                    preload_model_state = preload_model_state[opts["checkpoint_key"]]
+                if opts.get("prefix", ""):
+                    # Only params with this prefix should be loaded.
+                    # They are expected to be in the checkpoint without this prefix.
+                    # By adding the prefix to all params,
+                    # we make sure that only params matching this condition are loaded.
+                    # This is in line with the behavior of the TF engine.
+                    preload_model_state = {opts["prefix"] + key: value for key, value in preload_model_state.items()}
+                ignore_params = opts.get("ignore_params", [])
+                ignore_params_prefixes = opts.get("ignore_params_prefixes", [])
+                for key in list(preload_model_state.keys()):
+                    if key in ignore_params or any(
+                        [key.startswith(ignore_key) for ignore_key in ignore_params_prefixes]
+                    ):
+                        print(f"Ignoring variable {key}", file=log.v3)
+                        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)
+                missing_keys, _ = self._model.load_state_dict(preload_model_state, strict=False)
+                if not opts.get("ignore_missing", False):
+                    prefix_keys = [key for key in self._model.state_dict() if key.startswith(opts.get("prefix", ""))]
+                    missing_prefix_keys = set(prefix_keys).intersection(set(missing_keys))
+                    assert not missing_prefix_keys, f"Missing keys and ignore_missing=False: {missing_prefix_keys}"
+                print(f"Missing keys: {missing_keys}", file=log.v4)
 
         self._model.to(self._device)
 

{returnn-1.20230403.124714 → returnn-1.20230403.211148}/returnn/torch/frontend/_rand.py

@@ -10,7 +10,8 @@ import warnings
 
 def no_grad_trunc_normal_(tensor: torch.Tensor, mean, std, a, b, *, generator=None):
     """
-    Code copied and adopted from torch.nn.init._no_grad_trunc_normal_.
+    Code copied and adopted from torch.nn.init._no_grad_trunc_normal_,
+    to support the extra `generator` argument (https://github.com/pytorch/pytorch/issues/98200).
 
     Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
 

{returnn-1.20230403.124714 → returnn-1.20230403.211148/returnn.egg-info}/PKG-INFO

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

returnn-1.20230403.124714/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230403.124714'
-long_version = '1.20230403.124714+git.deb7a53'