returnn 1.20240828.151821__tar.gz → 1.20240829.92949__tar.gz

{returnn-1.20240828.151821 → returnn-1.20240829.92949}/PKG-INFO

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

returnn-1.20240829.92949/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240829.092949'
+long_version = '1.20240829.092949+git.448cffe'

{returnn-1.20240828.151821 → returnn-1.20240829.92949}/returnn/datasets/postprocessing.py

@@ -4,7 +4,8 @@ Provides :class:`PostprocessingDataset`.
 
 from __future__ import annotations
 
-import numpy as np
+from itertools import islice
+from numpy.random import Generator, PCG64
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from returnn.datasets.basic import DatasetSeq
@@ -98,7 +99,7 @@ class PostprocessingDataset(CachedDataset2):
         self._map_seq = map_seq
         self._map_seq_stream = map_seq_stream
         self._map_outputs = map_outputs
-        self._rng = np.random.default_rng(self._get_random_seed_for_epoch(0))
+        self._rng = Generator(PCG64(self._get_random_seed_for_epoch(0)))
 
         self._dataset = init_dataset(self._dataset_def, parent_dataset=self)
         if self._map_seq_stream is None:
@@ -143,7 +144,7 @@ class PostprocessingDataset(CachedDataset2):
             self._num_seqs = 0
             return True
 
-        self._rng = np.random.default_rng(self._get_random_seed_for_epoch(epoch=epoch))
+        self._rng = Generator(PCG64(self._get_random_seed_for_epoch(epoch=epoch)))
         assert self._dataset is not None
         self._dataset.init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
         self._data_iter = enumerate(self._build_mapping_iter())
@@ -180,7 +181,7 @@ class PostprocessingDataset(CachedDataset2):
         data_iter = self._iterate_dataset()
         if self._map_seq_stream is not None:
             data_iter = self._map_seq_stream(
-                data_iter, rng=self._rng, **{f"random-kwarg-{self._rng.integers(0, 1000)}": None}
+                data_iter, rng=self._rng, **{f"fwd_compatible_random_kwarg_{self._rng.integers(0, 1000)}": None}
             )
             assert isinstance(
                 data_iter, Iterator
@@ -201,7 +202,7 @@ class PostprocessingDataset(CachedDataset2):
                 tensor_dict.data[data_key].raw_tensor = self._dataset.get_data(seq_index, data_key)
             if self._map_seq is not None:
                 tensor_dict = self._map_seq(
-                    tensor_dict, rng=self._rng, **{f"random-kwarg-{self._rng.integers(0, 1000)}": None}
+                    tensor_dict, rng=self._rng, **{f"fwd_compatible_random_kwarg_{self._rng.integers(0, 1000)}": None}
                 )
                 assert isinstance(
                     tensor_dict, TensorDict
@@ -225,3 +226,40 @@ class PostprocessingDataset(CachedDataset2):
             if data_key in self._dataset.labels:
                 sparse_dim.vocab = Vocabulary.create_vocab_from_labels(self._dataset.labels[data_key])
         return Tensor(data_key, dims=dims, dtype=dtype, sparse_dim=sparse_dim)
+
+
+class LaplaceOrdering(Callable[[Iterator[TensorDict]], Iterator[TensorDict]]):
+    """
+    Iterator compatible with :class:`PostprocessingDataset`'s ``map_seq_stream`` applying
+    laplace sequence ordering based on the number of segments per bin.
+    """
+
+    def __init__(self, num_seqs_per_bin: int, length_key: str = "data"):
+        """
+        :param num_seqs_per_bin: number of segments in a single laplace bin.
+        :param length_key: data key to determine the segment length from for ordering.
+        """
+        self.length_key = length_key
+        assert num_seqs_per_bin > 0
+        self.num_seqs_per_bin = num_seqs_per_bin
+
+    def __call__(self, iterator: Iterator[TensorDict], **kwargs) -> Iterator[TensorDict]:
+        """:return: generator applying laplace sequence ordering on the data"""
+        iterator = iter(iterator)
+        is_down_phase = False
+
+        while True:
+            seq_buffer = list(islice(iterator, self.num_seqs_per_bin))
+            seq_buffer.sort(key=self._get_seq_len, reverse=is_down_phase)
+            yield from seq_buffer
+
+            is_down_phase = not is_down_phase
+
+            if len(seq_buffer) < self.num_seqs_per_bin:
+                break
+
+    def _get_seq_len(self, tdict: TensorDict) -> int:
+        """
+        :return: segment length of the segment in `tdict` as measured by `self.length_key` for comparison.
+        """
+        return tdict.data[self.length_key].raw_tensor.shape[0]

{returnn-1.20240828.151821 → returnn-1.20240829.92949}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240828.151821 → returnn-1.20240829.92949}/tests/test_Dataset.py

@@ -1087,6 +1087,30 @@ def test_PostprocessingDataset():
             assert len(classes) > 0
         assert count == 1
 
+    # test laplace ordering
+    with create_ogg_zip_txt_only_dataset_mult_seqs() as sub_ds_opts:
+        from returnn.datasets.postprocessing import LaplaceOrdering
+
+        ds_opts = {
+            "class": "PostprocessingDataset",
+            "dataset": sub_ds_opts,
+            "map_seq_stream": LaplaceOrdering(3, "classes"),
+        }
+        dataset = init_dataset(ds_opts)
+        dataset.init_seq_order(epoch=1)
+        assert dataset.have_seqs()
+        dataset.load_seqs(0, 6)
+
+        prev_len = None
+        for i in range(3):
+            classes = dataset.get_data(i, "classes")
+            assert prev_len is None or classes.shape[0] >= prev_len
+            prev_len = classes.shape[0]
+        for i in range(3, 6):
+            classes = dataset.get_data(i, "classes")
+            assert prev_len is None or classes.shape[0] <= prev_len or i == 3
+            prev_len = classes.shape[0]
+
 
 if __name__ == "__main__":
     better_exchook.install()

returnn-1.20240828.151821/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240828.151821'
-long_version = '1.20240828.151821+git.268f463'