returnn 1.20240905.105440__tar.gz → 1.20240905.152045__tar.gz

{returnn-1.20240905.105440 → returnn-1.20240905.152045}/PKG-INFO

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

returnn-1.20240905.152045/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240905.152045'
+long_version = '1.20240905.152045+git.c0f3741'

{returnn-1.20240905.105440 → returnn-1.20240905.152045}/returnn/datasets/distrib_files.py

@@ -417,8 +417,6 @@ class DistributeFilesDataset(CachedDataset2):
         return files_per_bin
 
     def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
-        if seq_idx >= self._num_seqs:
-            return None
         seq = self._workers[self.epoch].get_data_seq(seq_idx)
         if not seq:
             return None
@@ -429,11 +427,6 @@ class DistributeFilesDataset(CachedDataset2):
         """have seqs"""
         return bool(self.files)
 
-    @property
-    def num_seqs(self) -> int:
-        """num seqs"""
-        return self._num_seqs
-
     def finish_epoch(self, *, free_resources: bool = False):
         """finish epoch"""
         super().finish_epoch(free_resources=free_resources)
@@ -530,16 +523,18 @@ class _WorkerProcParent:
         # Thus, use the full epoch as the epoch here.
         # Make sure this init is async - we should not wait for this in the main process when it is not used yet.
         self.parent_conn.send(("init_seq_order", {"epoch": self.full_epoch_0idx + 1}))
-        self._num_seqs: Optional[int] = None
+        self._has_num_seqs = False
+        self._num_seqs: Optional[int] = None  # remains None if unknown
 
     def _lazy_wait_for_init_seq_order(self):
-        if self._num_seqs is not None:
+        if self._has_num_seqs:
             return
         msg, num_seqs = self.parent_conn.recv()
-        assert msg == "num_seqs" and isinstance(num_seqs, int)
+        assert msg == "num_seqs" and (num_seqs is None or isinstance(num_seqs, int))
         self._num_seqs = num_seqs
+        self._has_num_seqs = True
 
-    def get_num_seqs(self) -> int:
+    def get_num_seqs(self) -> Optional[int]:
         """num seqs for this sub epoch"""
         self._lazy_wait_for_init_seq_order()
         return self._num_seqs
@@ -656,7 +651,11 @@ def _worker_proc_loop(
             elif msg == "init_seq_order":
                 # We are responsible to get the seq order and distrib it to all the other workers.
                 dataset.init_seq_order(**kwargs)
-                parent_conn.send(("num_seqs", dataset.num_seqs))
+                try:
+                    num_seqs = dataset.num_seqs
+                except NotImplementedError:
+                    num_seqs = None
+                parent_conn.send(("num_seqs", num_seqs))
                 got_init_seq_order = True
                 next_seq_idx = 0
                 cache[:] = []

{returnn-1.20240905.105440 → returnn-1.20240905.152045}/returnn/datasets/postprocessing.py

@@ -173,8 +173,23 @@ class PostprocessingDataset(CachedDataset2):
         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())
+        if self._map_seq_stream is None:
+            # If we don't have an iterable mapper we know the number of segments exactly
+            # equals the number of segments in the wrapped dataset
+            try:
+                self._num_seqs = self._dataset.num_seqs
+            except NotImplementedError:
+                pass  # some datasets don't know their num_seqs
         return True
 
+    def get_data_keys(self):
+        """:return: available data keys"""
+        return list(self._out_tensor_dict_template.data.keys())
+
+    def get_data_dtype(self, key):
+        """:return: dtype of data entry `key`"""
+        return self._out_tensor_dict_template.data[key].dtype
+
     def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
         while True:
             try:

{returnn-1.20240905.105440 → returnn-1.20240905.152045}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240905.105440 → returnn-1.20240905.152045}/tests/test_Dataset.py

@@ -943,6 +943,10 @@ def test_DistributeFilesDataset_distribute_evenly_by_size():
     _test_stddev(rng.uniform(low=97, high=103, size=(100,)), 20, 0.03)
 
 
+def iter_identity(x, **kwargs):
+    yield from x
+
+
 def test_DistributeFilesDataset():
     from returnn.datasets.distrib_files import DistributeFilesDataset
     from test_HDFDataset import generate_hdf_from_other
@@ -1036,6 +1040,27 @@ def test_DistributeFilesDataset():
             global_seq_idx += 1
     assert global_seq_idx == num_hdf_files * num_seqs
 
+    # Test DistributeFilesDataset in conjunction w/ a dataset, that does not have a
+    # defined num_seqs:
+
+    def _get_subepoch_dset_w_unknown_num_seqs(files_subepoch: List[str]) -> Dict[str, Any]:
+        hdf_dset = {"class": "HDFDataset", "files": files_subepoch, "seq_ordering": "default"}
+        return {"class": "PostprocessingDataset", "dataset": hdf_dset, "map_seq_stream": iter_identity}
+
+    concat_dataset = init_dataset(
+        {
+            "class": "DistributeFilesDataset",
+            "files": hdf_files,
+            "get_sub_epoch_dataset": _get_subepoch_dset_w_unknown_num_seqs,
+            "partition_epoch": partition_epoch,
+        }
+    )
+    assert isinstance(concat_dataset, DistributeFilesDataset)
+    for sub_epoch in range(1, partition_epoch + 1):
+        print(f"Sub-epoch {sub_epoch}...")
+        concat_dataset.init_seq_order(epoch=sub_epoch)
+        concat_dataset.load_seqs(0, 1)
+
 
 def test_PostprocessingDataset():
     from returnn.tensor.tensor_dict import TensorDict
@@ -1053,6 +1078,7 @@ def test_PostprocessingDataset():
             "map_seq": _add_1337_to_classes,
         }
         dataset = init_dataset(ds_opts)
+        dataset.get_data_keys()
         dataset.init_seq_order(epoch=1)
         assert dataset.have_seqs()
         dataset.load_seqs(0, 1)

returnn-1.20240905.105440/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240905.105440'
-long_version = '1.20240905.105440+git.d79c168'