returnn 1.20241024.173653__tar.gz → 1.20241026.3853__tar.gz

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/PKG-INFO

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

returnn-1.20241026.3853/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241026.003853'
+long_version = '1.20241026.003853+git.98e9755'

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/returnn/__main__.py

@@ -576,6 +576,10 @@ def execute_main_task():
                 assert data, "set forward_data"
             else:
                 data = init_dataset(config.opt_typed_value("forward_data"))
+            # engine.epoch is usually the epoch of the loaded checkpoint,
+            # or what EngineBase.get_epoch_model will return.
+            # You can have both load and load_epoch, where load points to the checkpoint,
+            # and load_epoch is some other epoch, which you will get here for the dataset.
             data.init_seq_order(epoch=engine.epoch or 1)
             forward_callback = config.typed_value("forward_callback")
             assert forward_callback, "no forward_callback specified"

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/returnn/datasets/meta.py

@@ -1875,7 +1875,7 @@ class VariableDataset(Dataset):
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """init seq order"""
-        super().init_seq_order()
+        super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
         if epoch is None:
             if seq_list is not None or seq_order is not None:
                 raise ValueError(f"{self}: epoch is None, but given seq_list or seq_order, not supported")
@@ -1973,6 +1973,110 @@ class VariableDataset(Dataset):
         return self._dataset.is_data_sparse(key)
 
 
+class MultiEpochDataset(CachedDataset2):
+    """
+    It wraps some dataset, where one outer epoch corresponds to multiple epochs in the inner wrapped dataset.
+
+    This can be useful when the inner dataset uses partition_epoch, and we want to cover the whole full epoch.
+
+    One specific example when the data is distributed over multiple files,
+    and for reasonable performance, you want to have the data copied to the local disk,
+    but all data together is too large to fit on the local disk.
+    Then :class:`DistributeFilesDataset` is the logical choice,
+    which solves these issues.
+    However, you must use some partition_epoch in :class:`DistributeFilesDataset`
+    such that it will not load all data at once.
+    To cover all the data, you can use this :class:`MultiEpochDataset`
+    and set multi_epoch = partition_epoch of the inner dataset.
+    """
+
+    def __init__(self, *, dataset: Dict[str, Any], multi_epoch: int, **kwargs):
+        """
+        :param dataset: the inner wrapped dataset
+        :param multi_epoch: how much inner epochs correspond to one outer epoch
+        """
+        super().__init__(**kwargs)
+        self._dataset = init_dataset(dataset, parent_dataset=self)
+        assert self._dataset
+        self._multi_epoch = multi_epoch
+        self.num_inputs = self._dataset.num_inputs
+        self.num_outputs = self._dataset.num_outputs
+        self.labels = self._dataset.labels
+        self._cur_inner_start_epoch: Optional[int] = None
+        self._cur_inner_epoch_offset = 0
+        self._cur_inner_epoch_seq_idx_offset = 0
+        self._epoch_have_predefined_seq_order = False
+
+    def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
+        """init seq order"""
+        super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
+        self._epoch_have_predefined_seq_order = bool(seq_list or seq_order)
+        # epoch is 1-based
+        self._cur_inner_start_epoch = ((epoch - 1) * self._multi_epoch + 1) if epoch is not None else None
+        self._cur_inner_epoch_offset = 0
+        self._cur_inner_epoch_seq_idx_offset = 0
+        self._dataset.init_seq_order(epoch=self._cur_inner_start_epoch, seq_list=seq_list, seq_order=seq_order)
+
+    def finish_epoch(self, *, free_resources: bool = False):
+        """finish epoch"""
+        super().finish_epoch(free_resources=free_resources)
+        self._dataset.finish_epoch(free_resources=free_resources)
+
+    def get_all_tags(self) -> List[str]:
+        """all tags"""
+        return self._dataset.get_all_tags()
+
+    def get_total_num_seqs(self, *, fast: bool = False) -> int:
+        """total num seqs"""
+        return self._dataset.get_total_num_seqs(fast=fast)
+
+    def get_data_keys(self) -> List[str]:
+        """data keys"""
+        return self._dataset.get_data_keys()
+
+    def get_target_list(self) -> List[str]:
+        """target list"""
+        return self._dataset.get_target_list()
+
+    def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
+        assert seq_idx >= self._cur_inner_epoch_seq_idx_offset
+        sub_seq_idx = seq_idx - self._cur_inner_epoch_seq_idx_offset
+        if not self._dataset.is_less_than_num_seqs(sub_seq_idx):
+            if self._epoch_have_predefined_seq_order:
+                return None  # predefined seq order, so no multi-epoch handling
+            if self._cur_inner_start_epoch is None:
+                return None  # there was no epoch given, so no multi-epoch handling
+            self._cur_inner_epoch_offset += 1
+            if self._cur_inner_epoch_offset >= self._multi_epoch:
+                return None  # we are done
+            self._dataset.init_seq_order(epoch=self._cur_inner_start_epoch + self._cur_inner_epoch_offset)
+            self._cur_inner_epoch_seq_idx_offset = seq_idx
+            sub_seq_idx = 0
+            assert self._dataset.is_less_than_num_seqs(sub_seq_idx)  # expect that the sub epoch has some seqs
+        self._dataset.load_seqs(sub_seq_idx, sub_seq_idx + 1)
+        data = {}
+        for key in self.get_data_keys():
+            data[key] = self._dataset.get_data(sub_seq_idx, key)
+        seq_tag = self._dataset.get_tag(sub_seq_idx)
+        return DatasetSeq(seq_idx=seq_idx, seq_tag=seq_tag, features=data)
+
+    def get_data_dim(self, key: str) -> int:
+        """data dim"""
+        return self._dataset.get_data_dim(key)
+
+    def get_data_shape(self, data_key: str) -> List[int]:
+        """data shape"""
+        return self._dataset.get_data_shape(data_key)
+
+    def get_data_dtype(self, key: str) -> str:
+        """data dtype"""
+        return self._dataset.get_data_dtype(key)
+
+    def is_data_sparse(self, key: str) -> bool:
+        """is data sparse"""
+        return self._dataset.is_data_sparse(key)
+
+
 class AnythingDataset(Dataset):
     """
     An infinite dataset, creating dummy (zero) data on the fly,

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/returnn/engine/base.py

@@ -266,6 +266,12 @@ class EngineBase:
         """
         return self.pretrain and self.epoch == self.pretrain.get_train_num_epochs() + 1
 
+    def set_epoch(self, epoch: int):
+        """
+        Set the current epoch.
+        """
+        self.epoch = epoch
+
     def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIface):
         """
         Iterate through the dataset, calling `forward_step` from user config,

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/returnn/torch/engine.py

@@ -225,6 +225,11 @@ class Engine(EngineBase):
         self._train_step_func = self.config.typed_value("train_step")
         assert self._train_step_func, "train_step not defined"
 
+    def set_epoch(self, epoch: int):
+        """set epoch"""
+        super().set_epoch(epoch)
+        self._epoch_mp_shared.value = epoch
+
     def train(self):
         """
         Main training loop.
@@ -243,9 +248,7 @@ class Engine(EngineBase):
         )
         self.epoch = self._start_epoch - 1
         while self.epoch + 1 <= self._final_epoch:
-            self.epoch += 1
-            self._epoch_mp_shared.value = self.epoch
-
+            self.set_epoch(self.epoch + 1)
             self.init_train_epoch()
             self.train_epoch()
 

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/returnn.egg-info/PKG-INFO

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

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/tests/test_Dataset.py

@@ -18,12 +18,12 @@ from returnn.util.basic import NumbersDict
 from returnn.util import better_exchook
 
 
-def dummy_iter_dataset(dataset: Dataset) -> List[DatasetSeq]:
+def dummy_iter_dataset(dataset: Dataset, *, epoch: int = 1) -> List[DatasetSeq]:
     """
     :param Dataset dataset:
     :return: seqs
     """
-    dataset.init_seq_order(epoch=1)
+    dataset.init_seq_order(epoch=epoch)
     data_keys = dataset.get_data_keys()
     seq_idx = 0
     seqs = []
@@ -1147,6 +1147,67 @@ def test_PostprocessingDataset():
     assert func(2) == 21
 
 
+def test_MultiEpochDataset():
+    from returnn.datasets.meta import MultiEpochDataset
+    from returnn.datasets.cached2 import CachedDataset2
+
+    in_dim, out_dim = 11, 7
+    seq_len = 5
+    inner_num_seqs = 10
+
+    class _MyDataset(CachedDataset2):
+        def __init__(self):
+            super().__init__()
+            self.num_inputs = in_dim
+            self.num_outputs = {"classes": out_dim}
+
+        # noinspection PyShadowingNames
+        def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
+            """init seq order"""
+            super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
+            self._num_seqs = inner_num_seqs
+
+        def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
+            if seq_idx >= self._num_seqs:
+                return None
+            return DatasetSeq(
+                seq_idx=seq_idx,
+                seq_tag=repr({"epoch": self.epoch, "seq_idx": seq_idx}),
+                features=numpy.zeros((seq_len, in_dim)),
+                targets={"classes": numpy.zeros((seq_len,), dtype=numpy.int32)},
+            )
+
+    inner_dataset = _MyDataset()
+    inner_dataset.initialize()
+
+    multi_epoch = 3
+    dataset = MultiEpochDataset(dataset=inner_dataset, multi_epoch=multi_epoch)
+    for outer_epoch in [1, 7]:
+        seqs = dummy_iter_dataset(dataset, epoch=outer_epoch)
+        assert len(seqs) == inner_num_seqs * multi_epoch
+        outer_seq_idx = 0
+        sub_ep = (outer_epoch - 1) * multi_epoch + 1  # 1-based
+        sub_seq_idx = 0
+        for seq in seqs:
+            assert outer_seq_idx == seq.seq_idx
+            assert seq.features["data"].shape == (seq_len, in_dim)
+            assert seq.features["classes"].shape == (seq_len,)
+            print("seq:", seq.seq_tag)
+            d = eval(seq.seq_tag)  # seq tag is dict repr
+            assert isinstance(d, dict)
+            assert d["epoch"] == sub_ep
+            assert d["seq_idx"] == sub_seq_idx
+            # Calc next expected values.
+            if sub_seq_idx >= inner_num_seqs - 1:
+                sub_seq_idx = 0
+                sub_ep += 1
+            else:
+                sub_seq_idx += 1
+            outer_seq_idx += 1
+        assert outer_seq_idx == len(seqs)
+        assert sub_ep == outer_epoch * multi_epoch + 1 and sub_seq_idx == 0
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

{returnn-1.20241024.173653 → returnn-1.20241026.3853}/tests/test_torch_engine.py

@@ -294,6 +294,183 @@ def test_forward_beam_seq_lens():
         assert len(max_sizes) > 1
 
 
+def test_torch_engine_forward_dataset_epoch():
+    import tempfile
+    import shutil
+    import atexit
+    import os
+    import returnn
+
+    model_dir_name = tempfile.mkdtemp()
+    assert model_dir_name and os.path.isdir(model_dir_name) and not os.listdir(model_dir_name)
+    atexit.register(lambda: shutil.rmtree(model_dir_name))
+
+    in_dim, out_dim = 9, 13
+
+    def _get_model(**_kwargs):
+        return torch.nn.Linear(in_dim, out_dim)
+
+    epoch = 17
+    filename = Engine.epoch_model_filename(f"{model_dir_name}/model", epoch=epoch) + ".pt"
+
+    # That's how RETURNN now saves the model (2024-10-25).
+    # Maybe leave it like this for the test, even when RETURNN itself changes it,
+    # so that we also test that we still support this format.
+    torch.save(
+        {
+            "model": _get_model().state_dict(),  # some random model
+            "epoch": epoch,
+            "step": 123,
+            "effective_learning_rate": 0.13,
+            "returnn_version": returnn.__long_version__,
+        },
+        filename,
+    )
+
+    recent_seen_seq_idx: Optional[int] = None
+
+    class _ForwardCallback(ForwardCallbackIface):
+        def process_seq(self, *, seq_tag: str, outputs: TensorDict):
+            print("*** forward callback process seq", seq_tag)
+            d = eval(seq_tag)  # we prepared the dataset this way that we get some dict repr here...
+            assert isinstance(d, dict)
+            assert d["epoch"] == epoch
+            nonlocal recent_seen_seq_idx
+            seq_idx = d["seq_idx"]
+            if seq_idx == 0:
+                assert recent_seen_seq_idx is None
+            else:
+                assert recent_seen_seq_idx is not None
+                assert seq_idx == recent_seen_seq_idx + 1
+            recent_seen_seq_idx = seq_idx
+
+    forward_callback = _ForwardCallback()
+
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        print("*** forward step", extern_data)
+        data = extern_data["data"]
+        # Doesn't matter what we set as output here, not used...
+        # (Without output, maybe RETURNN complains, so put sth.)
+        # We just use the seq_tag in the forward callback, which is anyway available.
+        data.mark_as_default_output(shape=data.dims)
+
+    config = Config(
+        dict(
+            task="forward",
+            batch_size=50,
+            extern_data={"data": {"dim": in_dim}},
+            get_model=_get_model,
+            load=filename,
+            forward_step=_forward_step,
+            torch_dataloader_opts=dict(num_workers=0),  # simplifies the test
+        )
+    )
+
+    from returnn.datasets.cached2 import CachedDataset2
+    from returnn.datasets.basic import DatasetSeq
+
+    num_seqs = 10
+
+    class _MyDataset(CachedDataset2):
+        def __init__(self):
+            super().__init__()
+            self.num_inputs = in_dim
+            self.num_outputs = {"classes": out_dim}
+
+        # noinspection PyShadowingNames
+        def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
+            """init seq order"""
+            super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
+            self._num_seqs = num_seqs
+
+        def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
+            if seq_idx >= self._num_seqs:
+                return None
+            return DatasetSeq(
+                seq_idx=seq_idx,
+                seq_tag=repr({"epoch": self.epoch, "seq_idx": seq_idx}),
+                features=numpy.zeros((10, in_dim)),
+                targets={"classes": numpy.zeros((10,), dtype=numpy.int32)},
+            )
+
+    dataset = _MyDataset()
+    dataset.initialize()
+
+    with global_config_ctx(config):
+        engine = Engine(config=config)
+        engine.init_network_from_config()
+        # We expect that the engine epoch is set to the epoch of the checkpoint.
+        assert engine.epoch == epoch
+
+        for epoch in [3, 7, 11]:
+            engine.set_epoch(epoch)
+            assert engine.epoch == epoch
+            dataset.init_seq_order(epoch=epoch)
+            assert dataset.num_seqs == num_seqs
+            recent_seen_seq_idx = None
+            engine.forward_with_callback(callback=forward_callback, dataset=dataset)
+            assert recent_seen_seq_idx == num_seqs - 1
+
+
+def test_torch_engine_forward_load_epoch():
+    import tempfile
+    import shutil
+    import atexit
+    import os
+    import returnn
+
+    model_dir_name = tempfile.mkdtemp()
+    assert model_dir_name and os.path.isdir(model_dir_name) and not os.listdir(model_dir_name)
+    atexit.register(lambda: shutil.rmtree(model_dir_name))
+
+    in_dim, out_dim = 9, 13
+
+    def _get_model(**_kwargs):
+        return torch.nn.Linear(in_dim, out_dim)
+
+    epoch = 17
+    load_epoch = 11  # some other epoch
+    filename = Engine.epoch_model_filename(f"{model_dir_name}/model", epoch=epoch) + ".pt"
+
+    # That's how RETURNN now saves the model (2024-10-25).
+    # Maybe leave it like this for the test, even when RETURNN itself changes it,
+    # so that we also test that we still support this format.
+    torch.save(
+        {
+            "model": _get_model().state_dict(),  # some random model
+            "epoch": epoch,
+            "step": 123,
+            "effective_learning_rate": 0.13,
+            "returnn_version": returnn.__long_version__,
+        },
+        filename,
+    )
+
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        print("*** forward step", extern_data)
+        data = extern_data["data"]
+        data.mark_as_default_output(shape=data.dims)  # dummy...
+
+    config = Config(
+        dict(
+            task="forward",
+            batch_size=50,
+            extern_data={"data": {"dim": in_dim}},
+            get_model=_get_model,
+            load=filename,
+            load_epoch=load_epoch,
+            forward_step=_forward_step,
+            torch_dataloader_opts=dict(num_workers=0),  # simplifies the test
+        )
+    )
+
+    with global_config_ctx(config):
+        engine = Engine(config=config)
+        engine.init_network_from_config()
+        # We expect that even though we loaded the checkpoint, we now have the load_epoch.
+        assert engine.epoch == load_epoch
+
+
 def test_min_seq_len():
     from returnn.datasets.generating import DummyDataset
 

returnn-1.20241024.173653/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241024.173653'
-long_version = '1.20241024.173653+git.d07c658'