returnn 1.20230719.125702__tar.gz → 1.20230727.163234__tar.gz

{returnn-1.20230719.125702/returnn.egg-info → returnn-1.20230727.163234}/PKG-INFO

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

returnn-1.20230727.163234/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230727.163234'
+long_version = '1.20230727.163234+git.a7c14af'

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/demos/demo-rf.config

@@ -23,7 +23,7 @@ in_dim = Dim(9, name="in")
 out_dim = Dim(2, name="out")
 extern_data = {
     "data": {"dims": (batch_dim, time_dim, in_dim), "dtype": "float32"},
-    "classes": {"dims": (batch_dim, time_dim), "sparse_dim": out_dim, "dtype": "int32"},
+    "classes": {"dims": (batch_dim, time_dim), "sparse_dim": out_dim, "dtype": "int32", "available_for_inference": False},
 }
 model_outputs = {
     "output": {"dims": (batch_dim, time_dim, out_dim), "dtype": "float32"},

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/demos/demo-torch.config

@@ -24,7 +24,7 @@ num_inputs = 9
 num_outputs = 2
 extern_data = {
     "data": {"dim": num_inputs},
-    "classes": {"dim": num_outputs, "sparse": True},
+    "classes": {"dim": num_outputs, "sparse": True, "available_for_inference": False},
 }
 model_outputs = {
     "output": {"dim": num_outputs},
@@ -64,6 +64,7 @@ def forward_step(*, model: Model, extern_data: TensorDict):
     """
     data = extern_data["data"]
     out = model(data.raw_tensor)
+    rf.get_run_ctx().expected_outputs["output"].dims[1].dyn_size_ext.raw_tensor = data.dims[1].dyn_size_ext.raw_tensor
     rf.get_run_ctx().mark_as_default_output(tensor=out)
 
 

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/returnn/frontend/run_ctx.py

@@ -8,7 +8,6 @@ or forwarding loop.
 
 from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, Dict
-import logging
 from dataclasses import dataclass
 from returnn.tensor import Tensor, Dim, TensorDict
 import returnn.frontend as rf
@@ -400,16 +399,10 @@ def _output_tensor_from_raw(raw_tensor, *, dims: Sequence[Dim], name: str) -> Te
     # however, otherwise, it is ambiguous, so we print a warning that we always use the highest dim size.
     for axis, dim in enumerate(tensor.dims):
         if dim.dyn_size_ext and dim.dyn_size_ext.raw_tensor is None:
-            dim_value_raw = _backend.global_backend.get_shape_tuple_raw(raw_tensor)[axis]
-            dim_value = rf.convert_to_tensor(dim_value_raw, dims=())
-            dim_value = rf.cast(dim_value, dtype=dim.dyn_size_ext.dtype)
+            # Only non-scalar dyn sizes matter.
             if dim.dyn_size_ext.dims:
-                dyn_size = rf.full(dims=dim.dyn_size_ext.dims, fill_value=dim_value)
-                logging.warning(
+                raise Exception(
                     f"Output {name!r} {tensor}: Cannot infer dynamic size for dim {dim}. "
-                    f"Using {dyn_size} as fallback."
+                    f"You must explicitly specify the dyn size by assigning `{dim}.dyn_size_ext.raw_tensor = ...`."
                 )
-            else:
-                dyn_size = dim_value
-            dim.dyn_size_ext.raw_tensor = dyn_size.raw_tensor
     return tensor

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/returnn/tensor/tensor_dict.py

@@ -75,60 +75,87 @@ class TensorDict:
         return TensorDict({k: v.copy_template() for k, v in self.data.items()})
 
     def as_raw_tensor_dict(
-        self, *, include_const_sizes: bool = False, expected_type: Union[Type, Sequence[Type]] = object
+        self,
+        *,
+        include_const_sizes: bool = False,
+        include_scalar_dyn_sizes: bool = True,
+        exclude_duplicate_dims: bool = False,
+        expected_value_type: Union[Type, Sequence[Type]] = object,
     ) -> Dict[str, Any]:
         """
         :return: dict of raw tensors, including any sequence lengths / dynamic sizes
         """
+        assert not (include_const_sizes and not include_scalar_dyn_sizes)
+        visited_dims = set()
         out = {}
         for key, value in self.data.items():
             assert key not in out
             assert isinstance(
-                value.raw_tensor, expected_type
-            ), f"key {key} {value}: unexpected {type(value.raw_tensor)}, expected {expected_type}"
+                value.raw_tensor, expected_value_type
+            ), f"key {key} {value}: unexpected {type(value.raw_tensor)}, expected {expected_value_type}"
             out[key] = value.raw_tensor
             for i, dim in enumerate(value.dims):
+                if exclude_duplicate_dims and dim in visited_dims:
+                    continue
                 key_ = f"{key}:size{i}"
                 assert key_ not in out
                 if dim.is_batch_dim() and (not dim.dyn_size_ext or dim.dyn_size_ext.raw_tensor is None):
-                    dim_value = dim.get_dim_value()
-                    assert isinstance(
-                        dim_value, expected_type
-                    ), f"key {key_} {dim}: unexpected {type(dim_value)}, expected {expected_type}"
-                    out[key_] = dim_value
+                    if include_scalar_dyn_sizes:
+                        dim_value = dim.get_dim_value()
+                        assert isinstance(
+                            dim_value, expected_value_type
+                        ), f"key {key_} {dim}: unexpected {type(dim_value)}, expected {expected_value_type}"
+                        out[key_] = dim_value
                 elif dim.dyn_size_ext:
-                    assert isinstance(dim.dyn_size_ext.raw_tensor, expected_type), (
-                        f"key {key_} {dim} {dim.dyn_size_ext}:"
-                        f" unexpected {type(dim.dyn_size_ext.raw_tensor)}, expected {expected_type}"
-                    )
-                    out[key_] = dim.dyn_size_ext.raw_tensor
+                    if include_scalar_dyn_sizes or dim.dyn_size_ext.dims:
+                        assert isinstance(dim.dyn_size_ext.raw_tensor, expected_value_type), (
+                            f"key {key_} {dim} {dim.dyn_size_ext}:"
+                            f" unexpected {type(dim.dyn_size_ext.raw_tensor)}, expected {expected_value_type}"
+                        )
+                        out[key_] = dim.dyn_size_ext.raw_tensor
                 elif dim.size is not None:
-                    if include_const_sizes:
+                    if include_scalar_dyn_sizes and include_const_sizes:
                         assert isinstance(
-                            dim.size, expected_type
-                        ), f"key {key_} {dim}: unexpected {type(dim.size)}, expected {expected_type}"
+                            dim.size, expected_value_type
+                        ), f"key {key_} {dim}: unexpected {type(dim.size)}, expected {expected_value_type}"
                         out[key_] = dim.size
                 else:
                     raise Exception(f"cannot handle dim: {dim}")
+                visited_dims.add(dim)
         return out
 
-    def assign_from_raw_tensor_dict_(self, raw_tensor_dict: Dict[str, Any]):
+    def assign_from_raw_tensor_dict_(
+        self,
+        raw_tensor_dict: Dict[str, Any],
+        *,
+        with_scalar_dyn_sizes: bool = True,
+        duplicate_dims_are_excluded: bool = False,
+    ):
         """
         :param raw_tensor_dict: dict of raw tensors, including any sequence lengths / dynamic sizes
+        :param with_scalar_dyn_sizes: `include_scalar_dyn_sizes` was used in :func:`as_raw_tensor_dict`
+        :param duplicate_dims_are_excluded: `exclude_duplicate_dims` was used in :func:`as_raw_tensor_dict`
         """
+        visited_dims = set()
         for key, value in self.data.items():
             assert key in raw_tensor_dict
             value.raw_tensor = raw_tensor_dict[key]
             for i, dim in enumerate(value.dims):
+                if duplicate_dims_are_excluded and dim in visited_dims:
+                    continue
                 key_ = f"{key}:size{i}"
                 if dim.is_batch_dim() and not dim.dyn_size_ext:
                     dim.dyn_size_ext = Tensor("batch", [], dtype="int32")
                 if dim.dyn_size_ext:
-                    assert key_ in raw_tensor_dict
-                    dim.dyn_size_ext.raw_tensor = raw_tensor_dict[key_]
+                    if not with_scalar_dyn_sizes and not dim.dyn_size_ext.dims:
+                        pass
+                    else:
+                        assert key_ in raw_tensor_dict, f"keys: f{raw_tensor_dict.keys()}"
+                        dim.dyn_size_ext.raw_tensor = raw_tensor_dict[key_]
                 else:
                     if key_ in raw_tensor_dict:
                         assert dim.size == raw_tensor_dict[key_]
+                visited_dims.add(dim)
 
 
 def _convert_to_tensor(opts: _TensorT, *, name: Optional[str] = None) -> Tensor:

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/returnn/tf/native_op.py

@@ -267,7 +267,7 @@ class OpMaker(object):
       if(c->num_outputs() != %(num_outputs)i)
         return errors::InvalidArgument("wrong number of outputs. required %(num_outputs)i but got ", c->num_outputs());
       %(code_set_out_shape)s
-      return Status::OK();
+      return Status();
     })
     """ % {
             "num_inputs": len(in_info),

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/returnn/torch/engine.py

@@ -533,6 +533,8 @@ class Engine(EngineBase):
             raise TypeError(f"get_model returned {model} of type {type(model)}, expected rf.Module or torch.nn.Module")
         assert isinstance(self._pt_model, torch.nn.Module)
         print("Model:", self._pt_model, file=log.v4)
+        num_params = sum([parameter.numel() for parameter in self._pt_model.parameters()])
+        print(f"net params #: {num_params}", file=log.v2)
 
         if checkpoint_state is not None:
             self._pt_model.load_state_dict(checkpoint_state["model"])

{returnn-1.20230719.125702 → returnn-1.20230727.163234/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/tests/rf_utils.py

@@ -89,7 +89,7 @@ def run_model(
 
 def run_model_torch(extern_data: TensorDict, get_model: rf.GetModelFunc, forward_step: rf.StepFunc) -> TensorDict:
     """run"""
-    extern_data_raw = extern_data.as_raw_tensor_dict(expected_type=numpy.ndarray)
+    extern_data_raw = extern_data.as_raw_tensor_dict(expected_value_type=numpy.ndarray)
     rf.select_backend_torch()
     rf.set_random_seed(42)
 
@@ -111,7 +111,7 @@ def run_model_torch(extern_data: TensorDict, get_model: rf.GetModelFunc, forward
 
 def run_model_net_dict_tf(extern_data: TensorDict, get_model: rf.GetModelFunc, forward_step: rf.StepFunc) -> TensorDict:
     """run"""
-    extern_data_raw = extern_data.as_raw_tensor_dict(expected_type=numpy.ndarray)
+    extern_data_raw = extern_data.as_raw_tensor_dict(expected_value_type=numpy.ndarray)
     extern_data.reset_content()
     rf.select_backend_returnn_layers_tf()
     rf.set_random_seed(42)
@@ -159,9 +159,9 @@ def run_model_net_dict_tf(extern_data: TensorDict, get_model: rf.GetModelFunc, f
             layer = net.get_layer(layer_name)
             outputs_tf.data[k] = layer.output.copy()
 
-        fetches = outputs_tf.as_raw_tensor_dict(expected_type=tf.Tensor)
+        fetches = outputs_tf.as_raw_tensor_dict(expected_value_type=tf.Tensor)
         assert set(extern_data.data.keys()) == set(net.extern_data.data.keys())
-        extern_data_tf_placeholders = net.extern_data.as_raw_tensor_dict(expected_type=tf.Tensor)
+        extern_data_tf_placeholders = net.extern_data.as_raw_tensor_dict(expected_value_type=tf.Tensor)
         assert set(extern_data_tf_placeholders.keys()) == set(extern_data_raw.keys())
         feed_dict = {extern_data_tf_placeholders[k]: v for k, v in extern_data_raw.items()}
 

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/tests/test_demos.py

@@ -149,7 +149,13 @@ def test_demo_torch_task12ax():
     assert_less(fer, 0.02)
 
 
-def _test_torch_export_to_onnx(cfg_filename: str):
+def _test_torch_export_to_onnx(cfg_filename: str) -> str:
+    """
+    Executes the demo passed as a parameter and returns the ONNX exported model as a filename.
+
+    :param cfg_filename: Demo filename, either "demos/demo-rf.config" or "demos/demo-torch.config".
+    :return: Filename representing the ONNX model location.
+    """
     cleanup_tmp_models(cfg_filename)
     out = run(py, "rnn.py", cfg_filename, "--num_epochs", "1", "--device", "cpu")
     out_pt_model = re.search(r"\nSave model under (.*)\n", out, re.MULTILINE)
@@ -159,15 +165,69 @@ def _test_torch_export_to_onnx(cfg_filename: str):
     out_onnx_model = out_pt_model.replace(".pt", ".onnx")
     run(py, "tools/torch_export_to_onnx.py", cfg_filename, out_pt_model, out_onnx_model, print_stdout=True)
 
+    return out_onnx_model
+
+
+def _test_torch_onnx_inference_no_seq_lens(out_onnx_model: str):
+    """
+    Tests the inference of the torch demo with an ONNX model passed as parameter.
+    """
+    import onnxruntime as ort
+
+    torch.manual_seed(42)
+    # Set arbitrary values to pass to the model, respecting the input size (9).
+    dummy_data = torch.randn([3, 50, 9])
+
+    session = ort.InferenceSession(out_onnx_model)
+    outputs_onnx = session.run(
+        None,
+        {
+            "data": dummy_data.numpy(),
+        },
+    )
+    print("*** Result:", torch.FloatTensor(outputs_onnx[0]))
+    print("*** Batch size:", torch.IntTensor(outputs_onnx[1]))
+    print("*** Sequence lengths:", torch.IntTensor(outputs_onnx[2]))
+
+
+def _test_torch_onnx_inference_seq_lens_in_out(out_onnx_model: str):
+    """
+    Tests the inference of the torch demo with an ONNX model passed as parameter.
+    """
+    print(out_onnx_model)
+    import onnxruntime as ort
+
+    torch.manual_seed(42)
+    # Set arbitrary values to pass to the model, respecting the input size (9).
+    dummy_data = torch.randn([3, 50, 9])
+    dummy_seq_lens = torch.tensor([27, 50, 43], dtype=torch.int32)
+
+    session = ort.InferenceSession(out_onnx_model)
+    outputs_onnx = session.run(
+        None,
+        {
+            "data": dummy_data.numpy(),
+            "data:size1": dummy_seq_lens.numpy(),
+        },
+    )
+    out_result = torch.FloatTensor(outputs_onnx[0])
+    out_seq_lens = torch.IntTensor(outputs_onnx[1])
+    print("*** Result:", out_result)
+    print("*** Sequence lengths:", out_seq_lens)
+
+    assert out_result.shape == torch.Size([3, 50, 2])
+
 
 @unittest.skipIf(not torch, "no PyTorch")
 def test_demo_torch_export_to_onnx():
-    _test_torch_export_to_onnx("demos/demo-torch.config")
+    out_onnx_model = _test_torch_export_to_onnx("demos/demo-torch.config")
+    _test_torch_onnx_inference_seq_lens_in_out(out_onnx_model)
 
 
 @unittest.skipIf(not torch, "no PyTorch")
-def test_demo_rf_torch_export_to_onnx():
-    _test_torch_export_to_onnx("demos/demo-rf.config")
+def test_demo_rf_export_to_onnx():
+    out_onnx_model = _test_torch_export_to_onnx("demos/demo-rf.config")
+    _test_torch_onnx_inference_seq_lens_in_out(out_onnx_model)
 
 
 @unittest.skipIf(not torch, "no PyTorch")

{returnn-1.20230719.125702 → returnn-1.20230727.163234}/tools/torch_export_to_onnx.py

@@ -1,8 +1,33 @@
 """
-Converts a module from a config to ONNX. For that, it uses get_model() which must be available in the config
+Converts a RF/PT module to ONNX.
+
+For that, it uses `get_model(*, epoch, step, **_other_kwargs)` which must be available in the config
 and creates dummy data to be forwarded to the model.
 
-Since get_model() can return either a torch.nn.Module or a rf.Module, both cases must be taken into account.
+Since `get_model(...)` can return either a torch.nn.Module or a rf.Module, both cases are handled.
+
+`extern_data` defines the inputs.
+  - It strips away all entries with available_for_inference=False.
+  - Dynamic dimensions with dyn_size_ext being a scalar are excluded.
+  - If dynamic dims occur multiple times, only the first occurence is kept.
+    Note that the order depends on the order of the dict, as it is defined (not sorted alphabetically).
+
+`model_outputs` defines the output dimensions.
+  - Dynamic dimensions with dyn_size_ext being a scalar are excluded.
+
+For RASR, we usually expect that the input and output are always of the format [B, T, D] (batch, time, dim).
+We distinguish three cases on RASR side:
+
+- Input without sequence length are provided, output also without sequence lengths:
+    Define the input time dimension in `extern_data` to have `dyn_size_ext` being a scalar,
+    and same for `model_outputs`.
+
+- Input with sequence length are provided, output without sequence lengths,
+  meaning that RASR expects that the outputs have the same sequence length as the inputs:
+    Define the output time dimension in `model_outputs` to have `dyn_size_ext` being a scalar.
+
+- Input with sequence length are provided, output with sequence lengths:
+    The input and output time dimensions are as usual, having `dyn_size_ext` of shape [B].
 """
 
 from __future__ import annotations
@@ -83,10 +108,10 @@ class ForwardModulePT(torch.nn.Module):
         Wrapper to forward_step from the config.
         """
         extern_data = self.extern_data.copy_template()
-        extern_data.assign_from_raw_tensor_dict_(data)
+        extern_data.assign_from_raw_tensor_dict_(data, with_scalar_dyn_sizes=False, duplicate_dims_are_excluded=True)
         self.forward_step_func(model=self.model, extern_data=extern_data)
         _check_matching_outputs()
-        return rf.get_run_ctx().outputs.as_raw_tensor_dict()
+        return rf.get_run_ctx().outputs.as_raw_tensor_dict(include_scalar_dyn_sizes=False)
 
 
 class ForwardModuleRF(_RFModuleAsPTModule):
@@ -110,16 +135,16 @@ class ForwardModuleRF(_RFModuleAsPTModule):
         Wrapper to forward_step from the config.
         """
         extern_data = self.extern_data.copy_template()
-        extern_data.assign_from_raw_tensor_dict_(data)
+        extern_data.assign_from_raw_tensor_dict_(data, with_scalar_dyn_sizes=False, duplicate_dims_are_excluded=True)
         self.forward_step_func(model=self.rf_module, extern_data=extern_data)
         _check_matching_outputs()
-        return rf.get_run_ctx().outputs.as_raw_tensor_dict()
+        return rf.get_run_ctx().outputs.as_raw_tensor_dict(include_scalar_dyn_sizes=False)
 
 
 def _check_matching_outputs():
     rf.get_run_ctx().check_outputs_complete()
     model_outputs_raw_keys = set(_get_model_outputs_raw_keys())
-    outputs_raw = rf.get_run_ctx().outputs.as_raw_tensor_dict()
+    outputs_raw = rf.get_run_ctx().outputs.as_raw_tensor_dict(include_scalar_dyn_sizes=False)
     outputs_raw_keys = set(outputs_raw.keys())
     assert model_outputs_raw_keys == outputs_raw_keys, (
         f"Model outputs raw keys and output raw keys from forward_step don't match.\n"
@@ -138,7 +163,7 @@ def _get_model_outputs_raw_keys():
     for k, v in model_outputs.data.items():
         model_outputs_raw_keys.append(k)
         for i, dim in enumerate(v.dims):
-            if dim.is_batch_dim() or dim.is_dynamic():
+            if dim.dyn_size_ext and dim.dyn_size_ext.dims:
                 model_outputs_raw_keys.append(f"{k}:size{i}")
     return model_outputs_raw_keys
 
@@ -147,7 +172,7 @@ def main():
     """
     Main entry point
     """
-    parser = argparse.ArgumentParser(description="Converts a RF/PT module to ONNX.")
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
     parser.add_argument(
         "config",
         type=str,
@@ -193,10 +218,13 @@ def main():
     extern_data = TensorDict()
     extern_data.update(extern_data_dict, auto_convert=True)
     extern_data.reset_content()
+    for k, v in list(extern_data.data.items()):
+        if not v.available_for_inference:
+            del extern_data.data[k]
 
     tensor_dict_fill_random_numpy_(extern_data)
     tensor_dict_numpy_to_torch_(extern_data)
-    extern_data_raw = extern_data.as_raw_tensor_dict()
+    extern_data_raw = extern_data.as_raw_tensor_dict(include_scalar_dyn_sizes=False, exclude_duplicate_dims=True)
     model_outputs_raw_keys = _get_model_outputs_raw_keys()
 
     if is_pt_module:
@@ -214,6 +242,8 @@ def main():
     for k, v in list(extern_data.data.items()) + list(model_outputs.data.items()):
         dynamic_axes[k] = {i: dim.name for i, dim in enumerate(v.dims) if dim.is_dynamic() or dim.is_batch_dim()}
         for i, dim in enumerate(v.dims):
+            if dim.dyn_size_ext and dim.dyn_size_ext.dims == ():
+                continue
             if dim.dyn_size_ext:
                 dynamic_axes[f"{k}:size{i}"] = {
                     j: dim_.name
@@ -221,6 +251,10 @@ def main():
                     if dim_.is_dynamic() or dim_.is_batch_dim()
                 }
 
+    print("*** Input names:", list(extern_data_raw.keys()))
+    print("*** Output names:", model_outputs_raw_keys)
+    print("*** Dynamic axes:", dynamic_axes)
+
     export_func(
         pt_model_fwd,
         (extern_data_raw, {}),

returnn-1.20230719.125702/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230719.125702'
-long_version = '1.20230719.125702+git.1e8e66f'