returnn 1.20230522.95406__tar.gz → 1.20230524.133005__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230522.95406
+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.20230522.95406 → returnn-1.20230524.133005}/returnn/frontend/run_ctx.py

@@ -37,12 +37,12 @@ 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():
+def init_forward_step_run_ctx(expected_outputs: Optional[TensorDict] = None):
     """
     Call this at the beginning of a new forward step.
     """
     global _run_ctx
-    _run_ctx = RunCtx(stage="forward_step")
+    _run_ctx = RunCtx(stage="forward_step", expected_outputs=expected_outputs)
 
 
 def get_run_ctx() -> RunCtx:
@@ -67,7 +67,9 @@ class RunCtx:
     In forwarding, we expect that some output is being defined via mark_as_output().
     """
 
-    def __init__(self, *, stage: str, train_flag: Union[bool, Tensor] = False):
+    def __init__(
+        self, *, stage: str, train_flag: Union[bool, Tensor] = False, expected_outputs: Optional[TensorDict] = None
+    ):
         """
         :param stage:
             - "init"
@@ -78,6 +80,7 @@ class RunCtx:
         self._train_flag = train_flag
         self.losses = {}  # type: Dict[str, Loss]
         self.outputs = TensorDict()
+        self.expected_outputs = expected_outputs
 
     @property
     def stage(self) -> str:
@@ -177,20 +180,48 @@ 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
 
+        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:
         """
         Calls mark_as_output(tensor, "output", shape=shape).
@@ -203,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.20230522.95406 → returnn-1.20230524.133005/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230522.95406
+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.20230522.95406 → 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()
 
 
@@ -118,7 +120,9 @@ def main():
     """
     parser = argparse.ArgumentParser(description="Converts a RF/PT module to ONNX.")
     parser.add_argument(
-        "config", type=str, help="Filename to config file. Must have get_model(). Can optionally have export()."
+        "config",
+        type=str,
+        help="Filename to config file. Must have `get_model()` and `forward_step()`. Can optionally have `export()`.",
     )
     parser.add_argument("checkpoint", type=str, help="Checkpoint to RF module, considering the backend.")
     parser.add_argument("out_onnx_filename", type=str, help="Filename of the final ONNX model.")
@@ -127,7 +131,21 @@ def main():
     args = parser.parse_args()
 
     init(config_filename=args.config, checkpoint=args.checkpoint, log_verbosity=args.verbosity, device=args.device)
-    rf.init_forward_step_run_ctx()
+
+    model_outputs_dict = config.typed_value("model_outputs")
+    assert (
+        model_outputs_dict is not None
+    ), "The specified config needs to have explicit model outputs. Please define `model_outputs` in your config."
+    model_outputs = TensorDict()
+    model_outputs.update(model_outputs_dict, auto_convert=True)
+    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():
+                model_outputs_raw_keys.append(f"{k}:size{i}")
+
+    rf.init_forward_step_run_ctx(expected_outputs=model_outputs)
     rf.set_random_seed(42)
 
     get_model_func = config.typed_value("get_model")
@@ -158,23 +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()
-
-    model_outputs_dict = config.typed_value("model_outputs")
-    assert (
-        model_outputs_dict is not None
-    ), "The specified config needs to have explicit model outputs. Please define `model_outputs` in your config."
-    model_outputs = TensorDict()
-    model_outputs.update(model_outputs_dict, auto_convert=True)
-    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():
-                model_outputs_raw_keys.append(f"{k}:size{i}")
+    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.20230522.95406/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230522.095406'
-long_version = '1.20230522.095406+git.43fdbe0'