returnn 1.20241106.124322__tar.gz → 1.20241108.160059__tar.gz

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/PKG-INFO

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

returnn-1.20241108.160059/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241108.160059'
+long_version = '1.20241108.160059+git.492ddd8'

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/returnn/frontend/rand.py

@@ -69,13 +69,14 @@ __all__ = [
 
 def set_random_seed(seed: int):
     """
+    This initializes the random state of the backend
+    and also the step-based random state
+    (see :func:`get_static_step_based_seed`, only used when ``static=True`` in :func:`random`).
+
     Call this at the beginning of the program
     (after the RF backend was selected),
-    or when the model and computation graph is supposed to be reinitialized.
-
-    This initializes the random state of the backend and also the step-based random state.
-
-    This is *not* expected to be called after each epoch or step.
+    or when the model and computation graph is supposed to be reinitialized
+    or at the beginning of each epoch.
 
     :param seed: should depend on epoch or step
     """
@@ -124,6 +125,8 @@ def reset_step_random_state():
 
 def get_static_step_based_seed(*, size=None) -> Union[int, numpy.ndarray]:
     """
+    This is intended as a static seed for :func:`random` when ``static=True`` is used.
+
     :return: from the static step-based random state, get a seed
     """
     return _step_rnd.randint(2**31, size=size)
@@ -180,18 +183,20 @@ def random(
     :param int|float|Tensor|None bound: for uniform, defining the range [-bound, bound)
     :param int|float|Tensor|None minval: for uniform
     :param int|float|Tensor|None maxval: for uniform
-    :param int|list[int]|numpy.ndarray|None seed: If not given, uses self.network.random.randint,
-      i.e. then it is controlled by the global seed setting, and every layer would get its own seed.
-      If you specify it explicitly, make sure every :class:`RandomLayer` uses a different seed,
-      otherwise you would get the same random numbers everywhere.
+    :param int|list[int]|numpy.ndarray|None seed:
+        Only for the case ``static=True``.
+        If not given, uses self.network.random.randint,
+        i.e. then it is controlled by the global seed setting, and every layer would get its own seed.
+        If you specify it explicitly, make sure every :class:`RandomLayer` uses a different seed,
+        otherwise you would get the same random numbers everywhere.
     :param str|tf.random.Algorithm|None algorithm: see :class:`RandomStateInitLayer`
     :param Tensor|None explicit_state: You can pass the state explicitly here.
-      If not given, will be created automatically, and updated automatically.
-      You could pass a :class:`VariableLayer` with initial value via :class:`RandomStateInitLayer`,
-      or directly a :class:`RandomStateInitLayer`.
-      If auto_update_state is True, it must be a variable,
-      and every time a new random number is created, this variable is updated.
-      Otherwise (default), it will not be updated automatically.
+        If not given, will be created automatically, and updated automatically.
+        You could pass a :class:`VariableLayer` with initial value via :class:`RandomStateInitLayer`,
+        or directly a :class:`RandomStateInitLayer`.
+        If auto_update_state is True, it must be a variable,
+        and every time a new random number is created, this variable is updated.
+        Otherwise (default), it will not be updated automatically.
     :param bool|None auto_update_state: only used when you pass an explicit state
     :param bool|None static: if no state at all should be used. it just relies on the seed then.
     :param out: if given, will directly write into it, if possible by backend

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/returnn/torch/data/returnn_dataset_wrapper.py

@@ -45,7 +45,7 @@ class ReturnnDatasetResetMpSharedEpochCallback:
     def __call__(self):
         # dataset is likely a copy of the original dataset, either in the main process or in a worker process
         # Use epoch_mp_shared to get the current epoch correctly in worked processes
-        epoch = self.epoch_mp_shared.value
+        epoch = self.epoch_mp_shared.value or None
         self.dataset.init_seq_order(epoch=epoch)
 
 

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/returnn/torch/engine.py

@@ -34,7 +34,7 @@ from returnn.util import NumbersDict
 from returnn.util.basic import hms, NotSpecified
 from returnn.util.result_with_reason import ResultWithReason
 from returnn.util.debug import debug_shell
-from returnn.util.math import simplify_and_format_number
+from returnn.util.math import simplify_and_format_number, merge_random_seeds
 from returnn.forward_iface import ForwardCallbackIface
 
 from .updater import Updater
@@ -248,7 +248,7 @@ class Engine(EngineBase):
         print(
             f"Starting training at epoch {self._start_epoch}, global train step {self.global_train_step}", file=log.v3
         )
-        self.epoch = self._start_epoch - 1
+        self.set_epoch(self._start_epoch - 1)
         while self.epoch + 1 <= self._final_epoch:
             self.set_epoch(self.epoch + 1)
             self.init_train_epoch()
@@ -282,6 +282,17 @@ class Engine(EngineBase):
             }
         )
 
+        # Note: The RF/Torch default random number generator influences many things during training,
+        # such as dropout and other random operations inside the model,
+        # but also some potential shuffling in the dataset iterator.
+        # Also see Dataset._get_default_random_seed_offset() and Dataset._get_random_seed_for_epoch().
+        random_seed = self.config.int("random_seed", 42)
+        seed_data = [self.epoch, self.global_train_step, random_seed]
+        if self._torch_distributed_ctx:
+            seed_data.append(self._torch_distributed_ctx.rank())
+        random_seed = merge_random_seeds(seed_data)  # Join all seeds into one int.
+        rf.set_random_seed(random_seed)
+
     def _maybe_reset_dev_memory_caches(self, *, force: bool = False):
         if not force and not self._reset_dev_memory_caches:
             return
@@ -372,10 +383,9 @@ class Engine(EngineBase):
                 num_seqs_ = (
                     int(extern_data_raw["num_seqs"]) if extern_data_raw.get("num_seqs", None) is not None else -1
                 )
-                last_seq_idx_ = extern_data_raw["seq_idx"].max()
-                assert last_seq_idx_ >= last_seq_idx
-                last_seq_idx = int(last_seq_idx_)
-                del last_seq_idx_
+                # Note: The batches might have been shuffled,
+                # thus we cannot really assert that the seq_idx is always increasing.
+                last_seq_idx = max(int(extern_data_raw["seq_idx"].max()), last_seq_idx)
                 if step_idx == 0:
                     if num_seqs_ >= 0:
                         print(f"Epoch {self.epoch} num_seqs: {num_seqs_}", file=log.v5)
@@ -655,7 +665,10 @@ class Engine(EngineBase):
             for key, value in losses.items():
                 losses_[key] = value
                 if key in score_keys:
-                    losses_[f"{key}:exp"] = math.exp(value)
+                    try:
+                        losses_[f"{key}:exp"] = math.exp(value)
+                    except OverflowError:
+                        losses_[f"{key}:exp"] = float("inf")
             losses = NumbersDict(losses_)
         return losses
 
@@ -688,6 +701,32 @@ class Engine(EngineBase):
         max_seqs = self.config.typed_value("max_seqs", -1)
         batches_dataset = data_pipeline.BatchingIterDataPipe(wrapped_dataset, batch_size=batch_size, max_seqs=max_seqs)
 
+        online_shuffle_batches = self.config.typed_value("online_shuffle_batches", None)
+        if train and online_shuffle_batches:
+            if isinstance(online_shuffle_batches, int):
+                online_shuffle_batches = {"buffer_size": online_shuffle_batches}
+            elif isinstance(online_shuffle_batches, dict):
+                if "buffer_size" not in online_shuffle_batches:
+                    raise ValueError(
+                        f"config online_shuffle_batches, buffer_size not defined, got {online_shuffle_batches}"
+                    )
+            else:
+                raise TypeError(
+                    f"config online_shuffle_batches, expected int or dict, got {type(online_shuffle_batches)}"
+                )
+            # Note on random seed: This is handled by the PyTorch DataLoader iterator logic and IterDataPipe reset.
+            # Specifically, when we create a new DataLoader iterator,
+            # this will get fetch a new random number (from current Torch RNG state),
+            # use that as seed for the shuffle buffer.
+            # Note: In case of distributed training, it will broadcast the seed from rank 0 to all others.
+            # This is maybe not really what we want?
+            # https://discuss.pytorch.org/t/shuffleriterdatapipe-but-different-random-seed-per-distributed-rank/212612
+            # I currently don't really see a good way to override this behavior.
+            # 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)
+
         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)}"
 
@@ -927,7 +966,7 @@ class Engine(EngineBase):
             # We decremented the step above.
             epoch += 1
             step += 1
-        self.epoch = epoch  # in training, this will be reset to start_epoch
+        self.set_epoch(epoch)  # in training, this will be reset to start_epoch
         self.global_train_step = step
 
         load_model_post_hooks = self.config.typed_value("load_model_post_hooks")
@@ -1237,7 +1276,7 @@ class Engine(EngineBase):
                 if not self._is_dataset_evaluated(name=name, epoch=epoch):
                     raise Exception(f"Scores of epoch {epoch} for {name} are missing.")
 
-        self.epoch = self._start_epoch - 1
+        self.set_epoch(self._start_epoch - 1)
         if not self._is_dataset_evaluated(name="train"):
             raise Exception(f"Scores of last train epoch {self.epoch} are missing.")
         for name in self.eval_datasets.keys():

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/returnn/util/math.py

@@ -3,8 +3,9 @@ Some mathematical functions, in pure NumPy.
 """
 
 from __future__ import annotations
-from typing import Union, Optional, Dict
+from typing import Union, Optional, Sequence, Dict
 import numpy
+import hashlib
 
 
 def ceil_div(a: int, b: int) -> int:
@@ -85,3 +86,19 @@ def simplify_and_format_number(n: Union[int, float]) -> str:
         return str(n).rstrip("0").rstrip(".")
     else:
         raise TypeError(f"Expected int or float, got {n!r} type {type(n)}")
+
+
+def merge_random_seeds(data_sources: Sequence[int], *, num_bytes: int = 4, signed: bool = False) -> int:
+    """
+    :param data_sources: A list of integers. We expect that they are all representable as 64-bit signed integers.
+    :param num_bytes: for the output seed.
+    :param signed: whether the output seed should be signed.
+    :return: A num_bytes*8-bit integer seed, deterministically derived from the input data.
+    """
+    # Convert each integer to bytes and concatenate them
+    combined = b"".join(int(source).to_bytes(8, "big", signed=True) for source in data_sources)
+    # Use SHA-256 to hash the combined bytes
+    hash_digest = hashlib.sha256(combined).digest()
+    # Convert the hash digest to an integer seed
+    seed = int.from_bytes(hash_digest[:num_bytes], "big", signed=signed)
+    return seed

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/returnn.egg-info/PKG-INFO

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

{returnn-1.20241106.124322 → returnn-1.20241108.160059}/tests/test_torch_engine.py

@@ -479,6 +479,7 @@ def test_min_seq_len():
     dataset.initialize()
     dataset.init_seq_order(epoch=1)
     engine = Engine(config=config)
+    engine.set_epoch(1)
     data_loader = engine._create_data_loader(dataset)
     for _ in data_loader:
         assert False, "Should not contain sequences"
@@ -488,6 +489,7 @@ def test_min_seq_len():
     dataset.initialize()
     dataset.init_seq_order(epoch=1)
     engine = Engine(config=config)
+    engine.set_epoch(1)
     data_loader = engine._create_data_loader(dataset)
     for _ in data_loader:
         return
@@ -502,6 +504,7 @@ def test_max_seq_len():
     dataset.initialize()
     dataset.init_seq_order(epoch=1)
     engine = Engine(config=config)
+    engine.set_epoch(1)
     data_loader = engine._create_data_loader(dataset)
     for _ in data_loader:
         assert False, "Should not contain sequences"
@@ -511,6 +514,7 @@ def test_max_seq_len():
     dataset.initialize()
     dataset.init_seq_order(epoch=1)
     engine = Engine(config=config)
+    engine.set_epoch(1)
     data_loader = engine._create_data_loader(dataset)
     for _ in data_loader:
         return
@@ -527,6 +531,7 @@ def test_data_loader_oggzip():
     with create_ogg_zip_txt_only_dataset_mult_seqs(num_seqs=ds_num_seqs, max_seq_len=ds_max_seq_len) as dataset:
         dataset.init_seq_order(epoch=1)
         engine = Engine(config=config)
+        engine.set_epoch(1)
         data_loader = engine._create_data_loader(dataset)
         num_batches = 0
         num_seqs = 0
@@ -551,6 +556,7 @@ def test_data_loader_oggzip():
     with create_ogg_zip_txt_only_dataset_mult_seqs(num_seqs=ds_num_seqs, max_seq_len=ds_max_seq_len) as dataset:
         dataset.init_seq_order(epoch=1)
         engine = Engine(config=config)
+        engine.set_epoch(1)
         data_loader = engine._create_data_loader(dataset)
         for batch in data_loader:
             assert isinstance(batch, dict)

returnn-1.20241106.124322/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241106.124322'
-long_version = '1.20241106.124322+git.bf0c605'