returnn 1.20230524.103608__tar.gz → 1.20230524.133005__tar.gz

{returnn-1.20230524.103608/returnn.egg-info → returnn-1.20230524.133005}/PKG-INFO

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

returnn-1.20230524.133005/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230524.133005'
+long_version = '1.20230524.133005+git.021392d'

{returnn-1.20230524.103608 → returnn-1.20230524.133005}/returnn/frontend/run_ctx.py

@@ -7,7 +7,7 @@ or forwarding loop.
 """
 
 from __future__ import annotations
-from typing import Optional, Union, Any, Sequence, Dict, List
+from typing import Optional, Union, Any, Sequence, Dict
 from dataclasses import dataclass
 from returnn.tensor import Tensor, Dim, TensorDict
 import returnn.frontend as rf
@@ -37,7 +37,7 @@ def init_train_step_run_ctx(*, train_flag: Union[bool, Tensor]):
     _run_ctx = RunCtx(stage="train_step", train_flag=train_flag)
 
 
-def init_forward_step_run_ctx(expected_outputs: Optional[List[str]] = None):
+def init_forward_step_run_ctx(expected_outputs: Optional[TensorDict] = None):
     """
     Call this at the beginning of a new forward step.
     """
@@ -68,7 +68,7 @@ class RunCtx:
     """
 
     def __init__(
-        self, *, stage: str, train_flag: Union[bool, Tensor] = False, expected_outputs: Optional[List[str]] = None
+        self, *, stage: str, train_flag: Union[bool, Tensor] = False, expected_outputs: Optional[TensorDict] = None
     ):
         """
         :param stage:
@@ -180,27 +180,47 @@ class RunCtx:
             If not specified, we try to infer BTF or BF as default, if that works, otherwise it will be an error.
         """
         assert self.stage == "forward_step"
+
+        if self.expected_outputs is not None:
+            assert (
+                name in self.expected_outputs.data
+            ), f"mark_as_output: unexpected output {name!r}, we expect outputs: {self.expected_outputs}"
+        expected_output = self.expected_outputs.data[name] if self.expected_outputs else None
+        if dims is None and expected_output:
+            dims = expected_output.dims
+        if dims is not None and expected_output:
+            assert (
+                expected_output.dims == dims
+            ), f"mark_as_output: {name!r} dims mismatch from expected output, given {dims}, expected {expected_output}"
+
         if not isinstance(tensor, Tensor):
             assert isinstance(tensor, _backend.global_backend.RawTensorType)
             tensor = rf.convert_to_tensor(tensor, dims=dims)
             # In case it was not specified, just accept whatever order we got.
             dims = tensor.dims
-        assert name not in self.outputs.data
+
         if dims is None:
             # We try some reasonable defaults, specifically: BTF or BF.
             dims = _default_dim_order(tensor)
         assert set(dims) == set(tensor.dims), f"mark_as_output: tensor {tensor} does not have the dims {dims}"
         tensor = tensor.copy_transpose(dims, allow_int=False)
         tensor = tensor.copy(name=name)
+        assert name not in self.outputs.data
         self.outputs.data[name] = tensor
 
-        # Sanity check: expected outputs should coincide with raw keys of output, including output sizes.
-        if self.expected_outputs is not None:
-            raw_outputs = self.outputs.as_raw_tensor_dict()
-            for k in self.expected_outputs:
-                assert (
-                    k in raw_outputs.keys()
-                ), f"Expected model output '{k}' not found in actual model output keys {list(raw_outputs.keys())}"
+        if expected_output:  # Sanity checks for expected output
+            assert expected_output.dims == tensor.dims, (
+                f"mark_as_output: {name!r} dims mismatch from expected output,"
+                f" given {dims}, expected {expected_output}"
+            )
+            assert expected_output.dtype == tensor.dtype, (
+                f"mark_as_output: {name!r} dtype mismatch from expected output,"
+                f" given {tensor.dtype}, expected {expected_output.dtype}"
+            )
+            assert expected_output.sparse_dim == tensor.sparse_dim, (
+                f"mark_as_output: {name!r} sparse_dim mismatch from expected output,"
+                f" given {tensor.sparse_dim}, expected {expected_output.sparse_dim}"
+            )
 
     def mark_as_default_output(self, tensor: Union[Tensor, Any], *, shape: Optional[Sequence[Dim]] = None) -> None:
         """
@@ -214,6 +234,17 @@ class RunCtx:
         """
         self.mark_as_output(tensor, "output", dims=shape)
 
+    def check_outputs_complete(self):
+        """
+        If expected outputs are given, check that all expected outputs are present.
+        """
+        if self.expected_outputs:
+            assert set(self.expected_outputs.data.keys()) == set(self.outputs.data.keys()), (
+                f"check_outputs_complete: expected outputs {self.expected_outputs} do not match"
+                f" actual outputs {self.outputs}"
+            )
+            # We don't need to check the dims, dtype, etc, as this is already done in mark_as_output.
+
     def total_loss(self) -> Union[Tensor, float]:
         """
         :return: total loss, as it is used for backpropagation

{returnn-1.20230524.103608 → returnn-1.20230524.133005/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230524.103608 → returnn-1.20230524.133005}/tools/torch_export_to_onnx.py

@@ -83,6 +83,7 @@ class ForwardModulePT(torch.nn.Module):
         extern_data = self.extern_data.copy_template()
         extern_data.assign_from_raw_tensor_dict_(data)
         self.forward_step_func(model=self.model, extern_data=extern_data)
+        rf.get_run_ctx().check_outputs_complete()
         return rf.get_run_ctx().outputs.as_raw_tensor_dict()
 
 
@@ -109,6 +110,7 @@ class ForwardModuleRF(_RFModuleAsPTModule):
         extern_data = self.extern_data.copy_template()
         extern_data.assign_from_raw_tensor_dict_(data)
         self.forward_step_func(model=self.rf_module, extern_data=extern_data)
+        rf.get_run_ctx().check_outputs_complete()
         return rf.get_run_ctx().outputs.as_raw_tensor_dict()
 
 
@@ -143,7 +145,7 @@ def main():
             if dim.is_batch_dim() or dim.is_dynamic():
                 model_outputs_raw_keys.append(f"{k}:size{i}")
 
-    rf.init_forward_step_run_ctx(expected_outputs=model_outputs_raw_keys)
+    rf.init_forward_step_run_ctx(expected_outputs=model_outputs)
     rf.set_random_seed(42)
 
     get_model_func = config.typed_value("get_model")
@@ -174,10 +176,12 @@ def main():
         model.load_state_dict(loaded_checkpoint["model"])
         model.eval()
         pt_model_fwd = ForwardModulePT(model, forward_step_func, extern_data)
-    else:
+    elif is_rf_module:
         pt_model_fwd = ForwardModuleRF(model, forward_step_func, extern_data)
         pt_model_fwd.load_state_dict(loaded_checkpoint["model"])
         pt_model_fwd.eval()
+    else:
+        assert False, "PT/RF module?"  # should not get here
 
     dynamic_axes = {}
     for k, v in list(extern_data.data.items()) + list(model_outputs.data.items()):

returnn-1.20230524.103608/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230524.103608'
-long_version = '1.20230524.103608+git.a6fdea8'