returnn 1.20250227.110407__py3-none-any.whl → 1.20250228.104237__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250227.110407
+Version: 1.20250228.104237
 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.20250227.110407'
-long_version = '1.20250227.110407+git.f069d06'
+version = '1.20250228.104237'
+long_version = '1.20250228.104237+git.bd23951'

returnn/torch/data/pipeline.py

@@ -337,7 +337,13 @@ class BucketOrderingIterDataPipe(torch.utils.data.IterDataPipe):
     """
 
     def __init__(
-        self, dataset: torch.utils.data.IterableDataset, *, buckets: Sequence[Tuple[int, int]], length_key: str
+        self,
+        dataset: torch.utils.data.IterableDataset,
+        *,
+        buckets: Sequence[Tuple[int, int]],
+        length_key: str,
+        random_bucket_prob: float = 0.0,
+        seed: Optional[int] = None,
     ):
         """
         :param dataset: dataset to apply bucket batching to
@@ -345,9 +351,17 @@ class BucketOrderingIterDataPipe(torch.utils.data.IterDataPipe):
             Segments longer than the largest size limit configured in the buckets are dropped. To avoid dropping
             any segments make sure your largest bucket allows segments larger than your longest training segment.
         :param length_key: data key to take as length measure
+        :param random_bucket_prob: Probability of putting a segment not into the best-fitting bucket, but into
+            a randomly chosen still-fitting bucket.
+            This increases seq length variation within the buckets at the cost of slighly more padding.
+        :param seed: random seed
         """
         self._dataset = dataset
         self._length_key = length_key
+        assert random_bucket_prob >= 0.0
+        self._random_bucket_prob = random_bucket_prob
+        self._rng = numpy.random.RandomState()
+        self._seed = seed % (2**32) if seed is not None else None
 
         assert buckets, "empty bucket batching configuration"
         if not all(size > 0 and max_seqs > 0 for size, max_seqs in buckets):
@@ -367,6 +381,12 @@ class BucketOrderingIterDataPipe(torch.utils.data.IterDataPipe):
             if bucket_idx >= len(self._max_seq_lens):
                 # seg is too long, drop it
                 continue
+            if (
+                self._random_bucket_prob > 0.0
+                and bucket_idx < len(self._max_seq_lens) - 1
+                and self._rng.rand() < self._random_bucket_prob
+            ):
+                bucket_idx = self._rng.randint(bucket_idx, len(self._max_bucket_sizes))
             buckets[bucket_idx].append(data_dict)
             if len(buckets[bucket_idx]) >= self._max_bucket_sizes[bucket_idx]:
                 yield buckets[bucket_idx]
@@ -383,6 +403,21 @@ class BucketOrderingIterDataPipe(torch.utils.data.IterDataPipe):
     def __getitem__(self, index):
         raise Exception(f"{self.__class__.__name__}.__getitem__ is not supported")
 
+    def set_seed(self, seed: int) -> BucketOrderingIterDataPipe:
+        """
+        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"""
+        if self._seed is None:
+            self._seed = int(2**31 + torch.empty((), dtype=torch.int32).random_().item())
+        self._rng.seed(self._seed)
+        self._seed = None
+
 
 def get_batching_iterable_dataset_from_config(
     *, dataset: torch.utils.data.IterableDataset, config: Config, train: bool
@@ -497,7 +532,7 @@ class ShufflingDataPipe(torch.utils.data.IterDataPipe):
         self._buffer_size = buffer_size
         self._monotonic_data_keys = monotonic_data_keys
         self._rng = numpy.random.RandomState()
-        self._seed = seed
+        self._seed = seed % (2**32) if seed is not None else None
 
     def __iter__(self):
         # The implementation is very similar to the PostprocessingDataset's combinator LaplaceOrdering.
@@ -547,8 +582,10 @@ class ShufflingDataPipe(torch.utils.data.IterDataPipe):
 
     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._seed = int(2**31 + torch.empty((), dtype=torch.int32).random_().item())
         self._rng.seed(self._seed)
         self._seed = None
 

{returnn-1.20250227.110407.dist-info → returnn-1.20250228.104237.dist-info}/METADATA

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

{returnn-1.20250227.110407.dist-info → returnn-1.20250228.104237.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=7HDw-iONH2fjzNktp78WrdUg15TWTFaD-8VGJITrmF0,5215
+returnn/PKG-INFO,sha256=WyCpdpBjUpCVa51XYBBPnFFDZgy5pufGhH0PkInknok,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=OM77n2pKDjDW1JMMarYcFEKDu0Dw8yhE2fdmNGkW-Lk,77
+returnn/_setup_info_generated.py,sha256=AsveviXvDuDQsTAWwDCtvPC3oMGwJWCILWlydtbAeak,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
@@ -211,7 +211,7 @@ 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=mwbvocYe8dOxRpbGOyH6S1QO_egwfn9FGFG-NGGyOsA,27823
+returnn/torch/data/pipeline.py,sha256=mA6R1QU9vvRmfaUBvdqI9jQeIB3O-01ODcpmXs1SZ-w,29458
 returnn/torch/data/queued_data_iter.py,sha256=PoOsGHdHVZjTmcyfq_ZOw--P6hyfTdmAWIRGq_Z_nLM,888
 returnn/torch/data/returnn_dataset_wrapper.py,sha256=2CaDapzrlqahANuq-nyVAtv5ENHuM8A7okORwYJDisg,8006
 returnn/torch/data/tensor_utils.py,sha256=-Teqi--LLbt6q_5mDRdoHZHmPgSdC83W706ukif_YiU,1284
@@ -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.20250227.110407.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250227.110407.dist-info/METADATA,sha256=7HDw-iONH2fjzNktp78WrdUg15TWTFaD-8VGJITrmF0,5215
-returnn-1.20250227.110407.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-returnn-1.20250227.110407.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250227.110407.dist-info/RECORD,,
+returnn-1.20250228.104237.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250228.104237.dist-info/METADATA,sha256=WyCpdpBjUpCVa51XYBBPnFFDZgy5pufGhH0PkInknok,5215
+returnn-1.20250228.104237.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+returnn-1.20250228.104237.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250228.104237.dist-info/RECORD,,