returnn 1.20241118.95601__tar.gz → 1.20241120.12835__tar.gz

{returnn-1.20241118.95601 → returnn-1.20241120.12835}/PKG-INFO

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

returnn-1.20241120.12835/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241120.012835'
+long_version = '1.20241120.012835+git.8747efe'

{returnn-1.20241118.95601 → returnn-1.20241120.12835}/returnn/engine/base.py

@@ -28,6 +28,12 @@ class EngineBase:
         """
         if config is None:
             config = get_global_config(auto_create=True)
+        # In case the Engine is used directly by some user script,
+        # make sure that we are properly initialized.
+        if not log.initialized:
+            log.init_by_config(config)
+        if not util.BehaviorVersion.is_set():
+            util.BehaviorVersion.set(config.int("behavior_version", None))
         self.config = config
         self.epoch = 0
         self.global_train_step = None  # type: Optional[int]
@@ -272,11 +278,19 @@ class EngineBase:
         """
         self.epoch = epoch
 
-    def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIface):
+    def forward_with_callback(
+        self, *, dataset: Dataset, callback: ForwardCallbackIface, dataset_init_epoch: bool = True
+    ):
         """
         Iterate through the dataset, calling `forward_step` from user config,
         collecting outputs in `rf.get_run_ctx()` via `mark_as_output` calls,
         and then calling `callback` for each entry.
+
+        :param dataset:
+        :param callback: see :class:`ForwardCallbackIface`
+        :param dataset_init_epoch: whether the engine will call ``dataset.init_seq_order`` at the beginning,
+            using the current epoch.
+            If False, it assumes that ``dataset.init_seq_order`` was already called.
         """
         raise NotImplementedError
 

{returnn-1.20241118.95601 → returnn-1.20241120.12835}/returnn/tf/engine.py

@@ -35,7 +35,7 @@ from returnn.tf.updater import Updater
 from returnn.tf.data_pipeline import FeedDictDataProvider, DatasetDataProvider
 import returnn.tf.horovod as tf_horovod
 import returnn.util.basic as util
-from returnn.util.basic import hms, NumbersDict, BackendEngine, BehaviorVersion
+from returnn.util.basic import hms, NumbersDict, BackendEngine
 from returnn.forward_iface import ForwardCallbackIface
 from pprint import pprint
 
@@ -877,10 +877,6 @@ class Engine(EngineBase):
         :param returnn.config.Config|None config:
         """
         super(Engine, self).__init__(config=config)
-        if not log.initialized:
-            log.init_by_config(self.config)
-        if not BehaviorVersion.is_set():
-            BehaviorVersion.set(self.config.int("behavior_version", None))
         if BackendEngine.selected_engine is None:
             # This is only the case if the Engine is initialized directly in a custom script.
             # Otherwise, the RETURNN main entry point would already have selected the engine.
@@ -2356,7 +2352,9 @@ class Engine(EngineBase):
 
         writer.close()
 
-    def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIface):
+    def forward_with_callback(
+        self, *, dataset: Dataset, callback: ForwardCallbackIface, dataset_init_epoch: bool = True
+    ):
         """forward"""
         # https://github.com/rwth-i6/returnn/issues/1336
         raise NotImplementedError("TF engine does not support the generic forward func yet...")

{returnn-1.20241118.95601 → returnn-1.20241120.12835}/returnn/torch/data/returnn_dataset_wrapper.py

@@ -49,6 +49,15 @@ class ReturnnDatasetResetMpSharedEpochCallback:
         self.dataset.init_seq_order(epoch=epoch)
 
 
+class ReturnnDatasetResetNoOpCallback:
+    """
+    Can be used as reset_callback.
+    """
+
+    def __call__(self):
+        pass
+
+
 class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
     """
     Converts a RETURNN dataset into a PyTorch IterableDataset.

{returnn-1.20241118.95601 → returnn-1.20241120.12835}/returnn/torch/engine.py

@@ -672,16 +672,24 @@ class Engine(EngineBase):
             losses = NumbersDict(losses_)
         return losses
 
-    def _create_data_loader(self, dataset: Dataset, *, train: bool = False) -> DataLoader:
+    def _create_data_loader(
+        self, dataset: Dataset, *, train: bool = False, dataset_init_epoch: bool = True
+    ) -> DataLoader:
         """
         :param dataset: RETURNN dataset
+        :param train: Train might use a separate batch size in the config (batch_size_train vs batch_size_dev).
+            Also online_shuffle_batches is only used in training.
+        :param dataset_init_epoch: Whether to call dataset.init_seq_order(epoch=self.epoch) or not.
         :return: PyTorch data loader created from given RETURNN dataset
         """
         # Make sure that _dataset_reset does not keep a ref to `self`,
         # otherwise it would trigger to pickle `self` and all its members.
-        dataset_reset = returnn_dataset_wrapper.ReturnnDatasetResetMpSharedEpochCallback(
-            dataset=dataset, epoch_mp_shared=self._epoch_mp_shared
-        )
+        if dataset_init_epoch:
+            dataset_reset = returnn_dataset_wrapper.ReturnnDatasetResetMpSharedEpochCallback(
+                dataset=dataset, epoch_mp_shared=self._epoch_mp_shared
+            )
+        else:
+            dataset_reset = returnn_dataset_wrapper.ReturnnDatasetResetNoOpCallback()
 
         wrapped_dataset = returnn_dataset_wrapper.ReturnnDatasetIterDataPipe(dataset, reset_callback=dataset_reset)
         if (self._min_seq_length is not None) or (self._max_seq_length is not None):
@@ -730,6 +738,13 @@ class Engine(EngineBase):
         loader_opts = self.config.typed_value("torch_dataloader_opts") or {}
         assert isinstance(loader_opts, dict), f"config torch_dataloader_opts, expected dict, got {type(loader_opts)}"
 
+        if not dataset_init_epoch:
+            loader_opts = loader_opts.copy()
+            # We want to keep the current initialized seq order of the dataset.
+            # Multiprocessing does not quite work with this
+            # (the serialization of the dataset state does not cover the current seq order).
+            loader_opts["num_workers"] = 0
+
         data_loader = data_pipeline.create_data_loader_from_batches(batches_dataset, loader_opts)
 
         if data_loader.num_workers > 0:  # uses multi processing
@@ -815,6 +830,13 @@ class Engine(EngineBase):
 
         self._create_model(epoch=epoch, step=step)
 
+        non_critical_for_restore_set = set()
+        if isinstance(self._orig_model, rf.Module):
+            for key, param in self._orig_model.named_parameters():
+                assert isinstance(param, rf.Parameter)
+                if param.non_critical_for_restore:
+                    non_critical_for_restore_set.add(key)
+
         self._ignore_param_set.clear()
         loaded_state_keys = set()
         missing_keys = set()
@@ -944,6 +966,14 @@ class Engine(EngineBase):
         if self._ignore_param_set:
             util_module.convert_parameters_to_buffers(self._pt_model, self._ignore_param_set, persistent=False)
 
+        if missing_keys.intersection(non_critical_for_restore_set):
+            # We ignore missing keys which are non-critical for restore.
+            print(
+                "Ignoring missing keys which are non-critical for restore:",
+                missing_keys.intersection(non_critical_for_restore_set),
+                file=log.v4,
+            )
+            missing_keys.difference_update(non_critical_for_restore_set)
         if missing_keys:
             raise Exception(
                 "\n".join(
@@ -1092,7 +1122,9 @@ class Engine(EngineBase):
             if os.path.isfile(filename):
                 os.unlink(filename)
 
-    def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIface):
+    def forward_with_callback(
+        self, *, dataset: Dataset, callback: ForwardCallbackIface, dataset_init_epoch: bool = True
+    ):
         """forward"""
         assert isinstance(dataset, Dataset)
         assert isinstance(callback, ForwardCallbackIface)
@@ -1119,7 +1151,7 @@ class Engine(EngineBase):
             f" we want to keep all source sentences."
         )
 
-        data_loader = self._create_data_loader(dataset)
+        data_loader = self._create_data_loader(dataset, dataset_init_epoch=dataset_init_epoch)
         if self._forward_auto_split_batch_on_oom:
             data_loader = QueuedDataIter(data_loader)
 

{returnn-1.20241118.95601 → returnn-1.20241120.12835}/returnn.egg-info/PKG-INFO

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

returnn-1.20241118.95601/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241118.095601'
-long_version = '1.20241118.095601+git.a0d992a'