returnn 1.20250226.183415__py3-none-any.whl → 1.20250228.101938__py3-none-any.whl

returnn/PKG-INFO

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

returnn/_setup_info_generated.py

@@ -1,2 +1,2 @@
-version = '1.20250226.183415'
-long_version = '1.20250226.183415+git.ba9d72e'
+version = '1.20250228.101938'
+long_version = '1.20250228.101938+git.c053cfd'

returnn/datasets/basic.py

@@ -16,6 +16,7 @@ from threading import RLock
 from random import Random, random
 import sys
 import os
+import math
 import numpy
 import functools
 import typing
@@ -937,28 +938,51 @@ class Dataset:
         else:
             # We don't know. So:
             # Some monotonic increasing function in [0,1] which never reaches 1.
-            import math
+            return max(1.0e-10, (1 - 1 / ((seq_idx**0.5) / 100 + 1)) * 0.99)
 
-            return max(1.0e-10, 1.0 - math.exp(-seq_idx * 1000))
-
-    def get_complete_frac(self, seq_idx):
+    def get_complete_frac(self, sorted_seq_idx: int, *, allow_only_lr_suitable: bool = False) -> Optional[float]:
         """
-        :param int seq_idx:
-        :return: Returns a fraction (float in [0,1], always > 0) of how far we have advanced
-          for this seq in the dataset.
-          This does not have to be exact. This is only for the user.
-        :rtype: float
+        Tries to calculate exactly how much of the current epoch is completed when
+        having processed seq ``sorted_seq_idx``.
+
+        ``sorted_seq_idx`` cannot be less than the seq index of the previously loaded seqs.
+
+        :param sorted_seq_idx: sorted seq idx
+        :param allow_only_lr_suitable: only return a value when that value is suitable/accurate enough
+            to base LR scheduling on it. If false, this function will return an approximative value
+            when the exact value cannot be calculated (due to unknown ``num_seqs``).
+            Approximative values can be appropriate for e.g. progress bars.
+        :return: continuous value in (0, 1] which represents how much of the current epoch
+            is completed after ``sorted_seq_idx``.
+            If ``allow_only_lr_suitable=True``, returns ``None`` if the value cannot be calculated such
+            that it is accurate enough for LR scheduling, and otherwises bases ``epoch_continuous`` on it
+            for any dynamic learning rate scheduling.
+            As ``sorted_seq_idx`` is monotonic, the return value is also guaranteed to be monotonic.
         """
         # noinspection PyBroadException
         try:
             num_seqs = self.num_seqs
         except Exception:  # num_seqs not always available
+            if allow_only_lr_suitable:
+                return None
+
             # noinspection PyBroadException
             try:
                 num_seqs = self.estimated_num_seqs
             except Exception:  # also not always available
                 num_seqs = None  # ignore
-        return self.generic_complete_frac(seq_idx, num_seqs)
+
+        if math.isinf(num_seqs):
+            if allow_only_lr_suitable:
+                # cannot compute meaningful complete_frac for infinite num_seqs
+                return None
+            else:
+                num_seqs = None
+
+        assert (
+            num_seqs is None or 0 <= sorted_seq_idx < num_seqs
+        ), f"{self}: invalid seq indices: 0 <= seq_idx ({sorted_seq_idx}) < num_seqs ({num_seqs}) violated"
+        return self.generic_complete_frac(sorted_seq_idx, num_seqs)
 
     @property
     def num_seqs(self) -> int:
@@ -1375,16 +1399,27 @@ class DatasetSeq:
     Encapsulates all data for one sequence.
     """
 
-    def __init__(self, seq_idx, features, targets=None, seq_tag=None):
+    def __init__(
+        self,
+        seq_idx: int,
+        features,
+        *,
+        targets=None,
+        seq_tag: Optional[str] = None,
+        complete_frac: Optional[float] = None,
+    ):
         """
-        :param int seq_idx: sorted seq idx in the Dataset
+        :param seq_idx: sorted seq idx in the Dataset
         :param numpy.ndarray|dict[str,numpy.ndarray] features: format 2d (time,feature) (float)
         :param dict[str,numpy.ndarray]|numpy.ndarray|None targets: name -> format 1d (time) (idx of output-feature)
-        :param str seq_tag: sequence name / tag
+        :param seq_tag: sequence name / tag
+        :param complete_frac: continuous value in (0, 1] which represents how much of the current epoch
+            has been consumed when this seq is processed
         """
         assert isinstance(seq_idx, (int, numpy.integer))
         self.seq_idx = int(seq_idx)
         self.seq_tag = seq_tag or ("seq-%i" % seq_idx)
+        self.complete_frac = complete_frac
         if not isinstance(features, dict):
             assert isinstance(features, numpy.ndarray)
             features = {"data": features}

returnn/datasets/cached2.py

@@ -228,6 +228,15 @@ class CachedDataset2(Dataset):
             keys.remove("data")
         return keys
 
+    def get_complete_frac(self, sorted_seq_idx, **kwargs):
+        """
+        :return: fractional completion value for the given sorted_seq_idx
+        """
+        seq = self._get_seq(sorted_seq_idx)
+        if seq is not None and seq.complete_frac is not None:
+            return seq.complete_frac
+        return super().get_complete_frac(sorted_seq_idx, **kwargs)
+
     def is_data_sparse(self, key):
         """
         :param str key: e.g. "data" or "classes"

returnn/datasets/distrib_files.py

@@ -601,7 +601,8 @@ def _worker_proc_loop(
         dataset.load_seqs(next_seq_idx, next_seq_idx + 1)
         seq_tag = dataset.get_tag(next_seq_idx)
         features = {data_key: dataset.get_data(next_seq_idx, data_key) for data_key in dataset.get_data_keys()}
-        res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features)
+        complete_frac = dataset.get_complete_frac(next_seq_idx, allow_only_lr_suitable=True)
+        res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features, complete_frac=complete_frac)
         cache.append(res)
         next_seq_idx += 1
         return True

returnn/datasets/meta.py

@@ -554,6 +554,12 @@ class MetaDataset(CachedDataset2):
         """
         return self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]
 
+    def get_complete_frac(self, sorted_seq_idx: int, **kwargs) -> Optional[float]:
+        """
+        :param sorted_seq_idx:
+        """
+        return self.datasets[self.default_dataset_key].get_complete_frac(sorted_seq_idx, **kwargs)
+
     def get_data_keys(self) -> List[str]:
         """data keys"""
         return sorted(self.data_keys)

returnn/datasets/multi_proc.py

@@ -75,6 +75,7 @@ class MultiProcDataset(CachedDataset2):
         self._seq_order_proc_parent_conn = None  # type: Optional[mpConnection]
         self._seq_order_proc = None  # type: Optional[mp.Process]
         self._worker_procs = None  # type: Optional[List[mp.Process]]
+        self._cur_max_complete_frac: Optional[float] = None
 
         if _meta_info_cache:
             # This allows to skip the lazy init in self.initialize().
@@ -246,7 +247,8 @@ class MultiProcDataset(CachedDataset2):
             dataset.load_seqs(next_seq_idx, next_seq_idx + 1)
             seq_tag = dataset.get_tag(next_seq_idx)
             features = {data_key: dataset.get_data(next_seq_idx, data_key) for data_key in dataset.get_data_keys()}
-            res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features)
+            complete_frac = dataset.get_complete_frac(next_seq_idx, allow_only_lr_suitable=True)
+            res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features, complete_frac=complete_frac)
             cache.append(res)
             next_seq_idx += 1
             return True
@@ -403,6 +405,7 @@ class MultiProcDataset(CachedDataset2):
             return True
 
         self._lazy_init()
+        self._cur_max_complete_frac = 0.0
 
         if self._sharding_method == "dedicated":
             for worker_conn in self._worker_parent_conns:
@@ -441,6 +444,12 @@ class MultiProcDataset(CachedDataset2):
         if data is None:
             return None
         assert isinstance(data, DatasetSeq)
+        # The complete_frac values from the subprocesses are not necessarily monotonic
+        # due to rounding errors in the sharding and such.
+        # We therefore fix them up here. This is valid due to monotonicity of `seq_idx`.
+        max_comp_frac = max(data.complete_frac, self._cur_max_complete_frac)
+        data.complete_frac = max_comp_frac
+        self._cur_max_complete_frac = max_comp_frac
         data.seq_idx = seq_idx
         return data
 

returnn/datasets/numpy_dump.py

@@ -154,4 +154,4 @@ class NumpyDumpDataset(Dataset):
     def _add_cache_seq(self, seq_idx, features, targets):
         last_seq_idx = self._get_cache_last_seq_idx()
         assert seq_idx == last_seq_idx + 1
-        self.cached_seqs += [DatasetSeq(seq_idx, features, targets)]
+        self.cached_seqs += [DatasetSeq(seq_idx, features, targets=targets)]

returnn/datasets/postprocessing.py

@@ -5,6 +5,7 @@ Provides :class:`PostprocessingDataset`.
 from __future__ import annotations
 
 from itertools import islice
+import numpy
 from numpy.random import RandomState
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar
 
@@ -57,6 +58,14 @@ class PostprocessingDataset(CachedDataset2):
             },
         }
 
+    The postprocessor functions operate on ``TensorDict``s, which have entries for
+    all data keys in the underlying dataset.
+
+    There may also be additional "meta" entries in the tensor dicts, like ``complete_frac``
+    and ``seq_tag``.
+    These should be copied over in a manner that is reasonable for the use case at hand and
+    ensures forwards compatibility as well as reasonably possible.
+
     The dataset itself does not support its own seq ordering and relies on the wrapped
     dataset for seq ordering instead. Specifying a ``seq_ordering`` other than ``default``
     results in an error.
@@ -155,7 +164,10 @@ class PostprocessingDataset(CachedDataset2):
             self._out_tensor_dict_template = self._in_tensor_dict_template.copy_template()
             self.labels = self._dataset.labels.copy()
         # update only after _out_tensor_dict_template has been created from _in_tensor_dict_template
-        self._in_tensor_dict_template.update({"seq_tag": {"dims": (), "dtype": "string"}}, auto_convert=True)
+        self._in_tensor_dict_template.update(
+            {"complete_frac": {"dims": (), "dtype": "float32"}, "seq_tag": {"dims": (), "dtype": "string"}},
+            auto_convert=True,
+        )
         self.num_outputs = {
             k: (t.sparse_dim.size if t.sparse_dim else t.shape[-1] if len(t.shape) > 0 else 1, t.ndim)
             for k, t in self._out_tensor_dict_template.data.items()
@@ -222,6 +234,15 @@ class PostprocessingDataset(CachedDataset2):
         """:return: dtype of data entry `key`"""
         return self._out_tensor_dict_template.data[key].dtype
 
+    def get_total_num_seqs(self, *, fast=False):
+        """:return: total num seqs excluding partition_epoch"""
+        if self._map_seq_stream is not None:
+            raise util.OptionalNotImplementedError(
+                f"{self}: get_total_num_seqs not allowed when map_seq_stream is set."
+            )
+        assert self._dataset is not None
+        return self._dataset.get_total_num_seqs(fast=fast)
+
     def supports_sharding(self) -> bool:
         """:return: whether this dataset supports sharding"""
         assert self._dataset is not None
@@ -249,11 +270,12 @@ class PostprocessingDataset(CachedDataset2):
             assert loaded_seq_idx <= seq_idx, "_collect_single_seq must be done monotonically"
             if loaded_seq_idx != seq_idx:
                 continue
-            seq = DatasetSeq(
-                features={k: t.raw_tensor for k, t in tensor_dict.data.items() if k != "seq_tag"},
-                seq_idx=seq_idx,
-                seq_tag=str(tensor_dict["seq_tag"].raw_tensor),
+            complete_frac = (
+                float(tensor_dict.data["complete_frac"].raw_tensor) if "complete_frac" in tensor_dict.data else None
             )
+            seq_tag = str(tensor_dict.data["seq_tag"].raw_tensor) if "seq_tag" in tensor_dict.data else f"seq-{seq_idx}"
+            features = {k: t.raw_tensor for k, t in tensor_dict.data.items() if k not in ["complete_frac", "seq_tag"]}
+            seq = DatasetSeq(complete_frac=complete_frac, features=features, seq_idx=seq_idx, seq_tag=seq_tag)
             return seq
 
     def _build_mapping_iter(self) -> Iterator[TensorDict]:
@@ -262,8 +284,20 @@ class PostprocessingDataset(CachedDataset2):
         """
 
         def _validate_tensor_dict_iter(inner: Iterator[TensorDict]) -> Iterator[TensorDict]:
+            last_complete_frac = 0.0
             for t_dict in inner:
-                assert "seq_tag" in t_dict.data, "seq_tag dropped from TensorDict in postprocessing pipeline"
+                assert isinstance(t_dict, TensorDict), (
+                    f"postprocessing mapper function must produce a {TensorDict.__name__}, "
+                    f"but got a {type(t_dict).__name__}"
+                )
+                if "complete_frac" in t_dict.data:  # sanity check complete_frac
+                    complete_frac = float(t_dict.data["complete_frac"].raw_tensor)
+                    assert 0.0 <= complete_frac <= 1.0, f"complete_frac must be in [0, 1], but got {complete_frac}"
+                    assert complete_frac >= last_complete_frac, (
+                        "complete_frac must be monotonically increasing, "
+                        f"but got {complete_frac} after {last_complete_frac}"
+                    )
+                    last_complete_frac = complete_frac
                 for data_key, out_t in self._out_tensor_dict_template.data.items():
                     in_t = t_dict.data[data_key]
                     assert (
@@ -294,8 +328,14 @@ class PostprocessingDataset(CachedDataset2):
             tensor_dict = self._in_tensor_dict_template.copy_template()
             for data_key in data_keys:
                 tensor_dict.data[data_key].raw_tensor = self._dataset.get_data(seq_index, data_key)
-            seq_tag_tensor = str_to_numpy_array(self._dataset.get_tag(seq_index))
-            tensor_dict.data["seq_tag"].raw_tensor = seq_tag_tensor
+
+            complete_frac = self._dataset.get_complete_frac(seq_index, allow_only_lr_suitable=True)
+            comp_frac_raw_tensor = None
+            if complete_frac is not None:
+                comp_frac_raw_tensor = numpy.array(complete_frac, dtype=numpy.float32)
+                tensor_dict.data["complete_frac"].raw_tensor = comp_frac_raw_tensor
+            seq_tag_raw_tensor = str_to_numpy_array(self._dataset.get_tag(seq_index))
+            tensor_dict.data["seq_tag"].raw_tensor = seq_tag_raw_tensor
 
             if self._map_seq is not None:
                 tensor_dict = self._map_seq(
@@ -305,10 +345,16 @@ class PostprocessingDataset(CachedDataset2):
                     tensor_dict, TensorDict
                 ), f"map_seq must produce a {TensorDict.__name__}, but produced {type(tensor_dict).__name__}"
 
-                # Re-adding the seq tag here causes no harm in case it's dropped since we don't
-                # add/drop any segments w/ the non-iterator postprocessing function.
+                # Re-adding the seq_tag/complete_frac here causes no harm in case they are dropped
+                # since we don't add/drop any segments w/ the non-iterator postprocessing function.
+                if "complete_frac" not in tensor_dict.data and comp_frac_raw_tensor is not None:
+                    tensor_dict.data["complete_frac"] = Tensor(
+                        "complete_frac", dims=(), dtype="float32", raw_tensor=comp_frac_raw_tensor
+                    )
                 if "seq_tag" not in tensor_dict.data:
-                    tensor_dict.data["seq_tag"].raw_tensor = seq_tag_tensor
+                    tensor_dict.data["seq_tag"] = Tensor(
+                        "seq_tag", dims=(), dtype="string", raw_tensor=seq_tag_raw_tensor
+                    )
 
                 if self._seq_list_for_validation is not None:
                     seq_tag = self._seq_list_for_validation[seq_index]
@@ -366,7 +412,12 @@ class LaplaceOrdering(Callable[[Iterator[TensorDict]], Iterator[TensorDict]]):
         seq_buffer = list(islice(iterator, self.num_seqs_per_bin))
         has_ended = False
         while True:
+            # Make sure to not reorder the monotonically increasing values for complete_frac
+            # so that the trainer can calculate the appropriate learning rates.
+            complete_frac_values = [tdict.data["complete_frac"].raw_tensor for tdict in seq_buffer]
             seq_buffer.sort(key=self._get_seq_len, reverse=is_down_phase)
+            for sorted_item, comp_frac in zip(seq_buffer, complete_frac_values):
+                sorted_item.data["complete_frac"].raw_tensor = comp_frac
 
             next_seq_buffer = []
 

returnn/datasets/raw_wav.py

@@ -83,7 +83,7 @@ class RawWavDataset(CachedDataset2):
         inputFeatures = inputFeatures.astype(np.float32)
         if outputFeatures is not None:
             outputFeatures = outputFeatures.astype(np.float32)
-        return DatasetSeq(seq_idx, inputFeatures, outputFeatures)
+        return DatasetSeq(seq_idx, inputFeatures, targets=outputFeatures)
 
     def _get_num_outputs(self, num_outputs):
         """

returnn/datasets/sprint.py

@@ -504,7 +504,7 @@ class SprintDatasetBase(Dataset):
                     assert seq_idx + 1 == self.next_seq_to_be_added
                     self.cond.wait()
 
-            self.added_data += [DatasetSeq(seq_idx, features, targets, seq_tag=segment_name)]
+            self.added_data += [DatasetSeq(seq_idx, features, targets=targets, seq_tag=segment_name)]
             self.cond.notify_all()
             return seq_idx
 
@@ -588,7 +588,7 @@ class SprintDatasetBase(Dataset):
         """
         self._complete_frac = frac
 
-    def get_complete_frac(self, seq_idx):
+    def get_complete_frac(self, seq_idx, **kwargs):
         """
         :param int seq_idx:
         :rtype: float

returnn/datasets/stereo.py

@@ -349,7 +349,7 @@ class StereoHdfDataset(StereoDataset):
         elif targets.shape[1] == 1:
             targets = np.reshape(targets.astype(np.int32), (targets.shape[0],))
 
-        return DatasetSeq(seq_idx, inputFeatures, targets)
+        return DatasetSeq(seq_idx, inputFeatures, targets=targets)
 
     @staticmethod
     def _normalizeVector(v, mean, variance):
@@ -438,4 +438,4 @@ class DatasetWithTimeContext(StereoHdfDataset):
         targets = None
         if "classes" in originalSeq.get_data_keys():
             targets = originalSeq.get_data("classes")
-        return DatasetSeq(seq_idx, inputFeatures, targets)
+        return DatasetSeq(seq_idx, inputFeatures, targets=targets)

returnn/torch/data/pipeline.py

@@ -21,6 +21,7 @@ other PyTorch datasets more directly, including also HuggingFace datasets.
 
 from __future__ import annotations
 import bisect
+import itertools
 from typing import Optional, Any, Sequence, Tuple, Union, List, Dict, Callable
 import sys
 from copy import deepcopy
@@ -65,6 +66,9 @@ def collate_batch(batch: List[Dict[str, numpy.ndarray]]) -> Dict[str, Union[torc
         if key in ("num_seqs", "epoch"):
             res[key] = batch[0][key]  # it should always be the same
             continue
+        elif key == "complete_frac":
+            res[key] = max(sample[key] for sample in batch)
+            continue
         ls = [create_tensor(sample[key]) for sample in batch]
         if not ls:
             raise ValueError("batch is empty?")
@@ -122,7 +126,7 @@ class ChunkingIterDataPipe(torch.utils.data.IterDataPipe):
 
             if not chunking_data_keys:
                 chunking_data_keys = list(data_dict.keys())  # use all if not configured separately
-                chunking_data_key_black_list = ["seq_tag", "seq_idx", "num_seqs", "epoch"]
+                chunking_data_key_black_list = ["seq_tag", "seq_idx", "num_seqs", "epoch", "complete_frac"]
                 for key in chunking_data_key_black_list:
                     if key in chunking_data_keys:
                         chunking_data_keys.remove(key)
@@ -269,8 +273,15 @@ class BatchingIterDataPipe(torch.utils.data.IterDataPipe):
         epoch = int(data_dict["epoch"])
         seq_idx = int(data_dict["seq_idx"])
         num_seqs = int(data_dict["num_seqs"])  # >=1 if known, otherwise -1
-        epoch_continuous = (epoch - 1 + (seq_idx + 1) / num_seqs) if num_seqs > 0 else None
-        return {"epoch": epoch, "seq_idx": seq_idx, "epoch_continuous": epoch_continuous, **get_fwd_compat_kwargs()}
+        complete_frac = float(data_dict["complete_frac"])  # >= 0 if known, otherwise -1
+        epoch_continuous = (
+            epoch - 1 + complete_frac
+            if complete_frac >= 0.0
+            else (epoch - 1 + (seq_idx + 1) / num_seqs)
+            if num_seqs > 0
+            else None
+        )
+        return {"epoch": epoch, "epoch_continuous": epoch_continuous, "seq_idx": seq_idx, **get_fwd_compat_kwargs()}
 
     def __iter__(self):
         """
@@ -455,6 +466,125 @@ class LenFilterDataPipe(torch.utils.data.IterDataPipe):
         raise Exception(f"{self.__class__.__name__}.__getitem__ not supported")
 
 
+class ShufflingDataPipe(torch.utils.data.IterDataPipe):
+    """
+    Data pipe that is similar to ``torch.utils.data.datapipes.iter.Shuffler``,
+    but it will keep certain data keys of the batches in order while shuffling the rest.
+
+    Used for e.g. ``complete_frac`` and ``seq_idx``.
+    """
+
+    def __init__(
+        self,
+        dataset: torch.utils.data.IterableDataset,
+        *,
+        buffer_size: int,
+        monotonic_data_keys: Sequence[str],
+        seed: Optional[int] = None,
+    ):
+        """
+        :param dataset: batches dataset to shuffle
+        :param buffer_size: buffer size for shuffling
+        :param monotonic_data_keys: data keys that will be excluded from shuffling/keep their order
+        :param seed: random seed
+        """
+        super().__init__()
+
+        self._dataset = dataset
+        self._buffer: List[List[Dict[str, Any]]] = []
+        self._next_buffer: List[List[Dict[str, Any]]] = []
+        assert buffer_size > 0
+        self._buffer_size = buffer_size
+        self._monotonic_data_keys = monotonic_data_keys
+        self._rng = numpy.random.RandomState()
+        self._seed = seed
+
+    def __iter__(self):
+        # The implementation is very similar to the PostprocessingDataset's combinator LaplaceOrdering.
+
+        data_iter = iter(self._dataset)
+
+        self._buffer.extend(itertools.islice(data_iter, self._buffer_size))
+        has_ended = False
+
+        while True:
+            # Make sure to not reorder the monotonic values from self._monotonic_data_keys.
+            # These can contain things like complete_frac, which should be kept in order.
+            ordered_data = {
+                key: [data_dict[key] for batch in self._buffer for data_dict in batch]
+                for key in self._monotonic_data_keys
+            }
+            self._rng.shuffle(self._buffer)
+            for key in self._monotonic_data_keys:
+                data_dicts = [data_dict for batch in self._buffer for data_dict in batch]
+                assert len(data_dicts) == len(ordered_data[key])
+                for ordered_value, data_dict in zip(ordered_data[key], data_dicts):
+                    data_dict[key] = ordered_value
+
+            for item in self._buffer:
+                yield item
+
+                try:
+                    if not has_ended:
+                        self._next_buffer.append(next(data_iter))
+                except StopIteration:
+                    has_ended = True
+
+            if len(self._buffer) < self._buffer_size:
+                assert has_ended and not self._next_buffer
+                break
+
+            self._buffer.clear()
+            self._buffer, self._next_buffer = self._next_buffer, self._buffer
+
+    def set_seed(self, seed: int) -> ShufflingDataPipe:
+        """
+        Sets the seed for the next invocation of ``__iter__``, for compatibility with
+        ``torch.utils.data.graph_settings.apply_random_seed``.
+        """
+        self._seed = seed % (2**32)  # seed must be within [0, 2**32) for seeding RandomState
+        return self
+
+    def reset(self):
+        """resets the internal state of the data pipe"""
+        self._buffer.clear()
+        self._next_buffer.clear()
+        if self._seed is None:
+            self._seed = int(torch.empty((), dtype=torch.int32).random_().item())
+        self._rng.seed(self._seed)
+        self._seed = None
+
+    def __getstate__(self):
+        state = (
+            self._dataset,
+            self._buffer,
+            self._next_buffer,
+            self._buffer_size,
+            self._monotonic_data_keys,
+            self._rng.get_state(),
+            self._seed,
+        )
+        if torch.utils.data.IterDataPipe.getstate_hook is not None:
+            return torch.utils.data.IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        (
+            self._dataset,
+            self._buffer,
+            self._next_buffer,
+            self._buffer_size,
+            self._monotonic_data_keys,
+            rng_state,
+            self._seed,
+        ) = state
+        self._rng = numpy.random.RandomState()
+        self._rng.set_state(rng_state)
+
+    def __getitem__(self, index):
+        raise Exception(f"{self.__class__.__name__}.__getitem__ not supported")
+
+
 def create_data_loader_from_batches(
     batches_dataset: torch.utils.data.Dataset, loader_opts: Optional[Dict[str, Any]] = None
 ) -> torch.utils.data.DataLoader:

returnn/torch/data/returnn_dataset_wrapper.py

@@ -102,6 +102,7 @@ class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
 
         try:
             data_keys = self._dataset.get_data_keys()
+            last_complete_frac = -1
 
             seq_index = 0
             while self._dataset.is_less_than_num_seqs(seq_index):
@@ -109,11 +110,24 @@ class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
                 data = {data_key: self._dataset.get_data(seq_index, data_key) for data_key in data_keys}
                 data["seq_tag"] = str_to_numpy_array(self._dataset.get_tag(seq_index))
                 data["seq_idx"] = numpy.array(seq_index)
-                # It's slightly redundant to have num_seqs in each entry,
+
+                # It's slightly redundant to have the following data in each entry,
                 # but it's difficult to pass this back to the main proc otherwise.
-                data["num_seqs"] = num_seqs
-                # epoch is also redundant, but that's the cleanest/simplest way to pass it on to BatchingIterDataPipe.
                 data["epoch"] = epoch
+                data["num_seqs"] = num_seqs
+
+                complete_frac = self._dataset.get_complete_frac(seq_index, allow_only_lr_suitable=True)
+                if complete_frac is not None:
+                    assert 0.0 <= complete_frac <= 1.0, f"complete_frac must be in [0, 1], but got {complete_frac}"
+                    assert complete_frac >= last_complete_frac, (
+                        "complete_frac must be monotonically increasing, "
+                        f"but got {complete_frac} after {last_complete_frac}"
+                    )
+                else:
+                    complete_frac = -1
+                data["complete_frac"] = numpy.array(complete_frac, dtype=numpy.float32)
+                last_complete_frac = complete_frac
+
                 yield data
                 seq_index += 1
 

returnn/torch/engine.py

@@ -399,6 +399,7 @@ class Engine(EngineBase):
                     {k: int(util.prod(extern_data_raw[k].shape[:2])) for k in keys_w_seq_len},
                 )
 
+                complete_frac = float(extern_data_raw["complete_frac"])
                 num_seqs, last_seq_idx = _get_num_seqs_last_seq_idx(
                     report_prefix=report_prefix,
                     extern_data_raw=extern_data_raw,
@@ -406,7 +407,13 @@ class Engine(EngineBase):
                     prev_num_seqs=num_seqs,
                     prev_last_seq_idx=last_seq_idx,
                 )
-                epoch_continuous = (self.epoch - 1 + (last_seq_idx + 1) / num_seqs) if num_seqs is not None else None
+                epoch_continuous = (
+                    self.epoch - 1 + complete_frac
+                    if complete_frac >= 0.0
+                    else (self.epoch - 1 + (last_seq_idx + 1) / num_seqs)
+                    if num_seqs is not None
+                    else None
+                )
 
                 # clear the gradients when every gradient accumulation loop starts
                 if zero_grad_next_step:
@@ -777,7 +784,9 @@ class Engine(EngineBase):
             # Also note that we are likely using persistent multiprocessing data loader workers,
             # so calling torch.utils.data.graph_settings.apply_random_seed here in the main proc
             # will not have an effect then.
-            batches_dataset = torch.utils.data.datapipes.iter.Shuffler(batches_dataset, **online_shuffle_batches)
+            batches_dataset = data_pipeline.ShufflingDataPipe(
+                batches_dataset, monotonic_data_keys=("complete_frac", "seq_idx"), **online_shuffle_batches
+            )
 
         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)}"

{returnn-1.20250226.183415.dist-info → returnn-1.20250228.101938.dist-info}/METADATA

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

{returnn-1.20250226.183415.dist-info → returnn-1.20250228.101938.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=tasZ4y9DTXOoBq1n6RhxHj7GEEim3NIV3shYE_6qnzs,5215
+returnn/PKG-INFO,sha256=I8nJH2i19lJSp03bggFS1YlTbOT-yFLg8yanKsDGZEk,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=qBFbuB1yN3adgVM5pXt2-Yq9vorjRNchNPL8kDKx44M,31752
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=i9lO16SJCJurcbJrmKJjUX_VLD7LAXplYmS6TPYAzTI,77
+returnn/_setup_info_generated.py,sha256=J7wtu2Asd11qxOS3X2dv_AblIP3xvjshEattiHywgzQ,77
 returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
@@ -13,23 +13,23 @@ returnn/native_op.py,sha256=yqpE7SqBqXq77FCVnWMloUwadWlslEk-VzdK7FMpt_U,244411
 returnn/pretrain.py,sha256=MHiXJZqkQFmDVyaYsGpd_Acv20wxl7Pr6s6qJzAT2FI,22648
 returnn/datasets/__init__.py,sha256=PvDlfDOaaopIeUIt0OSvHD2eHZkdkyE-sjMXf35EH5U,390
 returnn/datasets/audio.py,sha256=Gmj7a08dnvYh7Z-G1TNapz42L50AIcDE9JeIZaO1s1M,23334
-returnn/datasets/basic.py,sha256=gLssy9J7nfwm1teWHEPoHPynWUWm1MBCpjqVUPZyZPA,70519
+returnn/datasets/basic.py,sha256=EhgyOv9bGHY08rCTQpt1HN_vW3djP5RwJuxtbp53neM,72300
 returnn/datasets/bundle_file.py,sha256=KQNrS1MSf-4_idlK0c0KFwON-f5sEK0sWU15WpoMYpE,2380
 returnn/datasets/cached.py,sha256=DIRdWrxBmsZG8O_9eVxBO5mcdo4f5KU-Xb-4wVz59Io,25418
-returnn/datasets/cached2.py,sha256=STojLL2Ivvd0xMfZRlYgzsHKlikYKL-caZCIDCgc_9g,11773
-returnn/datasets/distrib_files.py,sha256=kyqIQILDPAO2TXr39hjslmDxIAc3pkY1UOoj8nuiFXo,27534
+returnn/datasets/cached2.py,sha256=_6pza3IG68JexaExhj1ld3fP6pE7T-G804driJ9Z_qo,12141
+returnn/datasets/distrib_files.py,sha256=_UlcrnaU1rA9v6D3H3X4dPhcA--09fNeVnWs9VNo0yg,27656
 returnn/datasets/generating.py,sha256=e2-SXcax7xQ4fkVW_Q5MgOLP6KlB7EQXJi_v64gVAWI,99805
 returnn/datasets/hdf.py,sha256=shif0aQqWWNJ0b6YnycpPjIVNsxjLrA41Y66-_SluGI,66993
 returnn/datasets/lm.py,sha256=h0IHUbze87njKrcD5eT1FRxde7elIio05n-BWiqmjFE,98805
 returnn/datasets/map.py,sha256=kOBJVZmwDhLsOplzDNByIfa0NRSUaMo2Lsy36lBvxrM,10907
-returnn/datasets/meta.py,sha256=wHquywF1C7-YWhcSFSAdDNc0nEHRjE-ks7YIEuDFMIE,94731
-returnn/datasets/multi_proc.py,sha256=7kppiXGiel824HM3GvHegluIxtiNAHafm-e6qh6W7YU,21948
+returnn/datasets/meta.py,sha256=0wQzRzjShLSYNFoGo_MdR5IT8arxHr9gFjUlEqb2rbY,94969
+returnn/datasets/multi_proc.py,sha256=aVjsLt2qjHnHOrEYCgIPCwNYE-f1fiGP6eZ8NGAr3A4,22583
 returnn/datasets/normalization_data.py,sha256=wOHrbO3612uWXpzLHHxksDw0qeVmQ42w7byBL9QMh9Q,14618
-returnn/datasets/numpy_dump.py,sha256=c2Xgn8cfWxvRNCBMraMCRuHsbmjVQ05sISlaYWIRlKg,5150
-returnn/datasets/postprocessing.py,sha256=G9QiMP3Qr0RmA1PL6fCXOUfa2e_iPzZq_Nfx_u7SNiI,19980
-returnn/datasets/raw_wav.py,sha256=UyC4dUARb9QL0KOGhYdt96R2N_61JvFSvcyHMT8vMnw,9136
-returnn/datasets/sprint.py,sha256=_RS3IFlI5sgkLmvPqvSirWCi7-yxys_m-EY232ec8sM,55446
-returnn/datasets/stereo.py,sha256=0Df0Omm4T4r60GEFa6sEvZdgkm6keEw-qcvIO4BoJew,17617
+returnn/datasets/numpy_dump.py,sha256=wl8bKIKAlff2HPJPtuu5wBg3TLOf16d2wLVB4lLAwTM,5158
+returnn/datasets/postprocessing.py,sha256=Jkad_KHMesdPFFg9NKi7U3sbPw-RzxfUX_vOgJsI7p0,23075
+returnn/datasets/raw_wav.py,sha256=M7eTHp4CTtLQf3yPTiJY-mSJYgZNxkGV9IFN9J1dq_4,9144
+returnn/datasets/sprint.py,sha256=YhhdNbBTuL_HCc3asgK3o6vgq5h5nMPH5nBFvsuwVjA,55464
+returnn/datasets/stereo.py,sha256=PkowC91bZWihIYuIZgyGgPcNwgq5jBvyxxu1nER-VhM,17633
 returnn/datasets/text_dict.py,sha256=BPE73nh6-vtSLy3SiDf4dpFl9RJorE7oO6l5y2FU3MI,9965
 returnn/datasets/util/__init__.py,sha256=rEKhSD6fyhDiQF-x7dUQMwa29JZu72SDm7mYcCcLghY,52
 returnn/datasets/util/feature_extraction.py,sha256=axtXDb9wcNpOmyhmW3WJUj5xda29TKkKvOcGGvq7ExA,23923
@@ -207,13 +207,13 @@ returnn/tf/util/open_fst.py,sha256=sZRDw4TbxvhGqpGdUJWy1ebvlZm4_RPhygpRw9uLAOQ,1
 returnn/torch/README.md,sha256=jzJ2FpOHW02vxN69yKaV97C9LI-hmvjBglKfdZXIDdc,85
 returnn/torch/__init__.py,sha256=MHEUyNHB20Vy89uKAqZoj6FxJKF1Gq3HW-i6ra1pNcI,24
 returnn/torch/distributed.py,sha256=skFyutdVztxgTEk3HHJ8S83qRWbNpkNT8Tj16Ic0_hE,6981
-returnn/torch/engine.py,sha256=8BIpdcrpbJL9HrvCX-hISh-14zW9aSrHGvRWT9s0zOk,77103
+returnn/torch/engine.py,sha256=sU9A96icaj65uaEkX4i4aUK3IrB2S19_Fb9_sueB_JE,77426
 returnn/torch/updater.py,sha256=GqtBvZpElPVMm0lq84JPl4NVLFFETZAzAbR0rTomSao,28249
 returnn/torch/data/__init__.py,sha256=6cLNEi8KoGI12PF6akN7mI_mtjlx-0hcQAfMYoExwik,132
 returnn/torch/data/extern_data.py,sha256=_uT_9_gd5HIh1IoRsrebVG-nufSnb7fgC5jyU05GxJg,7580
-returnn/torch/data/pipeline.py,sha256=cIdSVjQHP9gihdfy4Pk2yu1-w572Qk8L2v26RL503qU,23266
+returnn/torch/data/pipeline.py,sha256=C0CAG_jk1oZwrPlW9WdRTxV9OvPztbqKjwKHnf3lhok,27886
 returnn/torch/data/queued_data_iter.py,sha256=PoOsGHdHVZjTmcyfq_ZOw--P6hyfTdmAWIRGq_Z_nLM,888
-returnn/torch/data/returnn_dataset_wrapper.py,sha256=1Bw82-Ge_8m_DSDXZNqQ3zGDic2HQlp6jysELL0NVK0,7369
+returnn/torch/data/returnn_dataset_wrapper.py,sha256=2CaDapzrlqahANuq-nyVAtv5ENHuM8A7okORwYJDisg,8006
 returnn/torch/data/tensor_utils.py,sha256=-Teqi--LLbt6q_5mDRdoHZHmPgSdC83W706ukif_YiU,1284
 returnn/torch/frontend/__init__.py,sha256=AA48HZnC17ASuKA0EWy8loZ-Bib_yUtqF4T1wYvjst4,62
 returnn/torch/frontend/_backend.py,sha256=TqyDWNP4XCvJNNGn8jyxaT8BOEjVE24QCUR3qsTIS3A,101242
@@ -253,8 +253,8 @@ returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
 returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20250226.183415.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250226.183415.dist-info/METADATA,sha256=tasZ4y9DTXOoBq1n6RhxHj7GEEim3NIV3shYE_6qnzs,5215
-returnn-1.20250226.183415.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-returnn-1.20250226.183415.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250226.183415.dist-info/RECORD,,
+returnn-1.20250228.101938.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250228.101938.dist-info/METADATA,sha256=I8nJH2i19lJSp03bggFS1YlTbOT-yFLg8yanKsDGZEk,5215
+returnn-1.20250228.101938.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+returnn-1.20250228.101938.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250228.101938.dist-info/RECORD,,