returnn 1.20230815.215448__tar.gz → 1.20230816.2116__tar.gz

{returnn-1.20230815.215448/returnn.egg-info → returnn-1.20230816.2116}/PKG-INFO

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

returnn-1.20230816.2116/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230816.002116'
+long_version = '1.20230816.002116+git.e1ca657'

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/returnn/frontend/run_ctx.py

@@ -29,20 +29,20 @@ def reset_run_ctx():
     _run_ctx = None
 
 
-def init_train_step_run_ctx(*, train_flag: Union[bool, Tensor]):
+def init_train_step_run_ctx(*, train_flag: Union[bool, Tensor], step: Union[int, Tensor]):
     """
     Call this at the beginning of a new train step.
     """
     global _run_ctx
-    _run_ctx = RunCtx(stage="train_step", train_flag=train_flag)
+    _run_ctx = RunCtx(stage="train_step", train_flag=train_flag, step=step)
 
 
-def init_forward_step_run_ctx(expected_outputs: Optional[TensorDict] = None):
+def init_forward_step_run_ctx(*, expected_outputs: Optional[TensorDict] = None, step: Union[int, Tensor]):
     """
     Call this at the beginning of a new forward step.
     """
     global _run_ctx
-    _run_ctx = RunCtx(stage="forward_step", expected_outputs=expected_outputs)
+    _run_ctx = RunCtx(stage="forward_step", expected_outputs=expected_outputs, step=step)
 
 
 def get_run_ctx() -> RunCtx:
@@ -68,7 +68,12 @@ class RunCtx:
     """
 
     def __init__(
-        self, *, stage: str, train_flag: Union[bool, Tensor] = False, expected_outputs: Optional[TensorDict] = None
+        self,
+        *,
+        stage: str,
+        train_flag: Union[bool, Tensor] = False,
+        step: Union[int, Tensor] = 0,
+        expected_outputs: Optional[TensorDict] = None,
     ):
         """
         :param stage:
@@ -78,6 +83,7 @@ class RunCtx:
         """
         self._stage = stage
         self._train_flag = train_flag
+        self._step = step
         self.losses = {}  # type: Dict[str, Loss]
         self.outputs = TensorDict()
         self.expected_outputs = expected_outputs
@@ -98,6 +104,14 @@ class RunCtx:
         """
         return self._train_flag
 
+    @property
+    def step(self) -> Union[int, Tensor]:
+        """
+        :return: global train step, starting with 0, not reset after an epoch, i.e. ignoring the epochs.
+            In a graph-based backend, this can be dynamic.
+        """
+        return self._step
+
     def mark_as_loss(
         self,
         loss: Union[Tensor, Any],

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/returnn/tf/frontend_layers/config_entry_points.py

@@ -47,15 +47,19 @@ def get_net_dict(
     model = get_model_func(epoch=epoch, step=step)
 
     task = config.value("task", None)
+    step_tensor = rfl.make_layer({"class": "global_train_step"}, name="global_train_step")
     if task in {"train", "eval"}:
-        rf.init_train_step_run_ctx(train_flag=rfl.make_layer({"class": "train_flag"}, name="train_flag"))
+        rf.init_train_step_run_ctx(
+            train_flag=rfl.make_layer({"class": "train_flag"}, name="train_flag"),
+            step=step_tensor,
+        )
         train_step_func = get_global_config().typed_value("train_step")
         train_step_func(
             model=model,
             extern_data=extern_data,
         )
     elif task in {"forward", "search"}:
-        rf.init_forward_step_run_ctx()
+        rf.init_forward_step_run_ctx(step=step_tensor)
         forward_step_func = get_global_config().typed_value("forward_step")
         forward_step_func(
             model=model,

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/returnn/tf/frontend_layers/layer.py

@@ -429,9 +429,13 @@ class Layer:
         # However, this can always be done via the `name_scope` option,
         # if the name would not match otherwise.
         queue = [self.root]  # type: List[Layer]
+        visited: Set[Layer] = set()
         mod_in_layer = {}  # type: Dict[Tuple[Layer, RefIdEq[rf.Module]], Layer]
         while queue:
             ctx = queue.pop(0)
+            if ctx in visited:  # can happen when parents are reassigned
+                continue
+            visited.add(ctx)
 
             assert not ctx.parent or ctx.parent.children[ctx.name] is ctx
             for child in ctx.children.values():

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/returnn/torch/engine.py

@@ -459,10 +459,12 @@ class Engine(EngineBase):
         """
         if train_func:
             assert self._train_step_func is not None
-            rf.init_train_step_run_ctx(train_flag=train_flag)
+            rf.init_train_step_run_ctx(train_flag=train_flag, step=self.global_train_step)
         else:
             assert self._forward_step_func is not None, "define forward_step in the config"
-            rf.init_forward_step_run_ctx(expected_outputs=self._forward_step_expected_outputs)
+            rf.init_forward_step_run_ctx(
+                expected_outputs=self._forward_step_expected_outputs, step=self.global_train_step
+            )
 
         from returnn.torch.distributed import ddp_train_forward_ctx
 

{returnn-1.20230815.215448 → returnn-1.20230816.2116/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/tests/rf_utils.py

@@ -99,7 +99,7 @@ def run_model_torch(extern_data: TensorDict, get_model: rf.GetModelFunc, forward
     tensor_dict_numpy_to_torch_(extern_data)
 
     model = get_model(epoch=1, step=0)
-    rf.init_forward_step_run_ctx()
+    rf.init_forward_step_run_ctx(step=0)
     forward_step(model=model, extern_data=extern_data)
     outputs = rf.get_run_ctx().outputs
     assert outputs.data

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/tests/test_rf_rec.py

@@ -285,7 +285,7 @@ def test_zoneout_lstm_tf_layers_vs_rf_pt():
         new_model.lstm.bias.raw_tensor.copy_(torch.from_numpy(tf_lstm_bias))
 
     print("*** Forward")
-    rf.init_train_step_run_ctx(train_flag=False)
+    rf.init_train_step_run_ctx(train_flag=False, step=0)
     extern_data.reset_content()
     extern_data.assign_from_raw_tensor_dict_(extern_data_numpy_raw_dict)
     tensor_dict_numpy_to_torch_(extern_data)

{returnn-1.20230815.215448 → returnn-1.20230816.2116}/tools/torch_export_to_onnx.py

@@ -192,13 +192,14 @@ def main():
     ), "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)
-    rf.init_forward_step_run_ctx(expected_outputs=model_outputs)
-    rf.set_random_seed(42)
 
     loaded_checkpoint = torch.load(args.checkpoint, map_location=torch.device(args.device))
     epoch = loaded_checkpoint["epoch"]
     step = loaded_checkpoint["step"]
 
+    rf.init_forward_step_run_ctx(expected_outputs=model_outputs, step=step)
+    rf.set_random_seed(42)
+
     get_model_func = config.typed_value("get_model")
     assert get_model_func, "get_model() isn't specified in the config passed as a parameter."
     sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}

returnn-1.20230815.215448/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230815.215448'
-long_version = '1.20230815.215448+git.fb82c05'