returnn 1.20250708.165746__tar.gz → 1.20250717.120243__tar.gz

{returnn-1.20250708.165746/returnn.egg-info → returnn-1.20250717.120243}/PKG-INFO

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

returnn-1.20250717.120243/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250717.120243'
+long_version = '1.20250717.120243+git.c980de6'

{returnn-1.20250708.165746 → returnn-1.20250717.120243}/returnn/datasets/distrib_files.py

@@ -174,17 +174,21 @@ class DistributeFilesDataset(CachedDataset2):
 
         self.distrib_shard_files = distrib_shard_files
         if distrib_shard_files:
-            assert self._num_shards == 1 and self._shard_index == 0, (  # ensure defaults are set
-                f"{self}: Cannot use both dataset-sharding via properties _num_shards and _shard index "
-                f"and {self.__class__.__name__}'s own sharding implementation based on the trainings rank and size."
-            )
             if _distrib_info:
-                # If we're in a child process `_get_rank_and_size()` no longer works,
+                # We're in a child process.
+                # `_get_rank_and_size()` no longer works,
                 # so we pass the info about the shards via a pickled property.
                 # See also Dataset.__reduce__.
-                self._shard_index = _distrib_info["_shard_index"]
-                self._num_shards = _distrib_info["_num_shards"]
+                # _num_shards and _shard_index are already set, so just check.
+                assert (
+                    self._shard_index == _distrib_info["_shard_index"]
+                    and self._num_shards == _distrib_info["_num_shards"]
+                )
             else:
+                assert self._num_shards == 1 and self._shard_index == 0, (  # ensure defaults are set
+                    f"{self}: Cannot use both dataset-sharding via properties _num_shards and _shard index "
+                    f"and {self.__class__.__name__}'s own sharding implementation based on the trainings rank and size."
+                )
                 self._shard_index, self._num_shards = _get_rank_and_size()
         assert 0 <= self._shard_index < self._num_shards
 
@@ -524,6 +528,8 @@ def _get_rank_and_size() -> Tuple[int, int]:
 
         ctx = returnn.tf.horovod.get_ctx(config=config)
         return ctx.rank(), ctx.size()
+    elif config.typed_value("__debug_dummy_distributed_rank_and_size") is not None:
+        return config.typed_value("__debug_dummy_distributed_rank_and_size")
     else:
         return 0, 1
 

{returnn-1.20250708.165746 → returnn-1.20250717.120243}/returnn/frontend/audio/specaugment.py

@@ -26,10 +26,10 @@ def specaugment(
     """
     SpecAugment, https://arxiv.org/abs/1904.08779
     """
-    if not feature_dim:
+    if feature_dim is None:
         assert x.feature_dim
         feature_dim = x.feature_dim
-    if not max_consecutive_feature_dims:
+    if max_consecutive_feature_dims is None:
         max_consecutive_feature_dims = feature_dim.dimension // 5
     if global_train_step_dependent:
         with rf.set_default_device_ctx("cpu"):

{returnn-1.20250708.165746 → returnn-1.20250717.120243/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250708.165746 → returnn-1.20250717.120243}/tests/test_Dataset.py

@@ -1120,9 +1120,63 @@ def test_DistributeFilesDataset():
     assert global_seq_idx == num_hdf_files * num_seqs
 
 
-def test_PostprocessingDataset():
-    from returnn.tensor.tensor_dict import TensorDict
+def _dfd_get_sub_epoch_dataset(files_subepoch: List[str]) -> Dict[str, Any]:
+    return {"class": "HDFDataset", "files": files_subepoch, "seq_ordering": "default"}
+
+
+def test_DistributeFilesDataset_sharding():
+    # https://github.com/rwth-i6/returnn/issues/1678
+    import pickle
+    from returnn.datasets.distrib_files import DistributeFilesDataset
+    from test_HDFDataset import generate_hdf_from_dummy
+    from returnn.config import global_config_ctx, Config
+
+    # Create a few HDF files such that we can easily verify the data later.
+    hdf_files = []
+    # num_seqs = 23
+    for hdf_idx in range(20):
+        hdf_files.append(generate_hdf_from_dummy())
+
+    # Test to load via DistributeFilesDataset.
 
+    # Setup some torch.distributed dummy environment.
+    distrib_size = 2
+    with global_config_ctx(Config({"__debug_dummy_distributed_rank_and_size": (1, distrib_size)})):
+        partition_epoch = 5
+        assert len(hdf_files) % partition_epoch == 0  # just for easier testing here
+        dataset = init_dataset(
+            {
+                "class": "DistributeFilesDataset",
+                "files": hdf_files,
+                "buffer_size": 100,
+                "distrib_shard_files": True,
+                "get_sub_epoch_dataset": _dfd_get_sub_epoch_dataset,
+                "partition_epoch": partition_epoch,
+                "seq_ordering": "random",
+            }
+        )
+        assert isinstance(dataset, DistributeFilesDataset)
+        assert dataset._num_shards == distrib_size and dataset._shard_index == 1
+        s = pickle.dumps(dataset)
+        dataset = pickle.loads(s)
+        assert isinstance(dataset, DistributeFilesDataset)
+        assert dataset._num_shards == distrib_size and dataset._shard_index == 1
+        global_seq_idx = 0
+        for sub_epoch in range(1, partition_epoch + 1):
+            print(f"Sub-epoch {sub_epoch}...")
+            dataset.init_seq_order(sub_epoch)
+            local_seq_idx = 0
+            while dataset.is_less_than_num_seqs(local_seq_idx):
+                print(f"Sub-epoch {sub_epoch}, seq {local_seq_idx} (global seq {global_seq_idx})...")
+                dataset.load_seqs(local_seq_idx, local_seq_idx + 1)
+                data = dataset.get_data(local_seq_idx, "classes")
+                assert data.ndim == 1 and data.shape[0] > 1
+                local_seq_idx += 1
+                global_seq_idx += 1
+        # assert global_seq_idx == len(hdf_files) * num_seqs // distrib_size  # TODO not sure...?
+
+
+def test_PostprocessingDataset():
     _demo_txt = "some utterance text that has a few words"
 
     def _add_1337_to_classes(tdict: TensorDict, **kwargs) -> TensorDict:

returnn-1.20250708.165746/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250708.165746'
-long_version = '1.20250708.165746+git.eff9959'