returnn 1.20240828.114639__tar.gz → 1.20240828.163303__tar.gz

{returnn-1.20240828.114639 → returnn-1.20240828.163303}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240828.114639
+Version: 1.20240828.163303
 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.163303/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240828.163303'
+long_version = '1.20240828.163303+git.bea6e2f'

{returnn-1.20240828.114639 → returnn-1.20240828.163303}/returnn/datasets/postprocessing.py

@@ -4,7 +4,9 @@ Provides :class:`PostprocessingDataset`.
 
 from __future__ import annotations
 
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+from itertools import islice
+import numpy as np
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from returnn.datasets.basic import DatasetSeq
 from returnn.datasets.util.vocabulary import Vocabulary
@@ -43,8 +45,10 @@ class PostprocessingDataset(CachedDataset2):
                 "files": ["/path/to/data.hdf"],
             },
             # one of them, but not both:
-            "map_seq": map_seq,  # (data: TensorDict) -> TensorDict
-            "map_seq_stream": map_seqs,  # (iter: Iterator[TensorDict]) -> Iterator[TensorDict]
+            # (data: TensorDict, *, rng: numpy.random.Generator, **kwargs) -> TensorDict
+            "map_seq": map_seq,
+            # (iter: Iterator[TensorDict], *, rng: numpy.random.Generator, **kwargs) -> Iterator[TensorDict]
+            "map_seq_stream": map_seqs,
             # only required when data shapes change wrt. the wrapped dataset:
             "map_outputs": {
                 "data": {"dims": [time_dim, new_data_dim]},
@@ -55,20 +59,26 @@ class PostprocessingDataset(CachedDataset2):
     def __init__(
         self,
         dataset: Dict[str, Any],
-        map_seq: Optional[Union[Callable[[TensorDict], TensorDict]]] = None,
-        map_seq_stream: Optional[Callable[[Iterator[TensorDict]], Iterator[TensorDict]]] = None,
+        map_seq: Optional[Callable] = None,
+        map_seq_stream: Optional[Callable] = None,
         map_outputs: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
         """
         :param dataset: inner dataset to be post-processed
         :param map_seq: post processor function operating on the single-segment level.
+            Signature: `(data: TensorDict, *, rng: numpy.random.Generator, **kwargs) -> TensorDict`
             To avoid confusion on the order of how the processing functions are applied to the data, only one of
             `map_seq` and `map_seq_stream` can be specified at a time.
+            To ensure forwards compatibility, the function must accept `**kwargs` as its last argument.
+            This is enforced by passing randomly named parameters at runtime.
         :param map_seq_stream: post processor function operating on the multiple segment level via an iterator.
             Allows merging multiple segments into one, or generating multiple output segments from one input segment.
+            Signature: `(iter: Iterator[TensorDict], *, rng: numpy.random.Generator, **kwargs) -> Iterator[TensorDict]`
             To avoid confusion on the order of how the processing functions are applied to the data, only one of
             `map_seq` and `map_seq_stream` can be specified at a time.
+            To ensure forwards compatibility, the function must accept `**kwargs` as its last argument.
+            This is enforced by passing randomly named parameters at runtime.
         :param map_outputs: Type and axis specification of the outputs of the mapping functions,
             like extern_data and model_outputs.
             To simplify the common case when no shapes change, this value can be left unspecified. The dataset then
@@ -89,6 +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._dataset = init_dataset(self._dataset_def, parent_dataset=self)
         if self._map_seq_stream is None:
@@ -133,6 +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))
         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())
@@ -168,7 +180,9 @@ class PostprocessingDataset(CachedDataset2):
 
         data_iter = self._iterate_dataset()
         if self._map_seq_stream is not None:
-            data_iter = self._map_seq_stream(data_iter)
+            data_iter = self._map_seq_stream(
+                data_iter, rng=self._rng, **{f"random-kwarg-{self._rng.integers(0, 1000)}": None}
+            )
             assert isinstance(
                 data_iter, Iterator
             ), f"map_seq_stream must produce an {Iterator.__name__}, but produced {type(data_iter).__name__}"
@@ -187,7 +201,9 @@ class PostprocessingDataset(CachedDataset2):
             for data_key in data_keys:
                 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)
+                tensor_dict = self._map_seq(
+                    tensor_dict, rng=self._rng, **{f"random-kwarg-{self._rng.integers(0, 1000)}": None}
+                )
                 assert isinstance(
                     tensor_dict, TensorDict
                 ), f"map_seq must produce a {TensorDict.__name__}, but produced {type(tensor_dict).__name__}"
@@ -210,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.114639 → returnn-1.20240828.163303}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240828.114639
+Version: 1.20240828.163303
 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.114639 → returnn-1.20240828.163303}/tests/test_Dataset.py

@@ -1042,7 +1042,7 @@ def test_PostprocessingDataset():
 
     _demo_txt = "some utterance text that has a few words"
 
-    def _add_1337_to_classes(tdict: TensorDict) -> TensorDict:
+    def _add_1337_to_classes(tdict: TensorDict, **kwargs) -> TensorDict:
         tdict.data["classes"] += 1337
         return tdict
 
@@ -1062,7 +1062,7 @@ def test_PostprocessingDataset():
 
     count = 0
 
-    def _repeat2(input_iter: Iterator[TensorDict]) -> Iterator[TensorDict]:
+    def _repeat2(input_iter: Iterator[TensorDict], **kwargs) -> Iterator[TensorDict]:
         nonlocal count
 
         for tdict in input_iter:
@@ -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.114639/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240828.114639'
-long_version = '1.20240828.114639+git.b243128'