take torch num_workers into account for sharding

Author: NeoLegends

returnn/datasets/basic.py

@@ -68,12 +68,6 @@ def set_or_remove(key, value):
         set_or_remove("min_chunk_size", config.opt_typed_value("min_chunk_size", 0) or None)
         set_or_remove("chunking_variance", config.float("chunking_variance", 0))
 
-        dd_cfg = config.typed_value("dataset_distribution", "random_seed_offset")
-        assert dd_cfg in ["random_seed_offset", "shard"]
-        shard_index, num_shards = Dataset._get_rank_and_size(config) if dd_cfg == "shard" else 0, 1
-        set_or_remove("num_shards", num_shards)
-        set_or_remove("shard_index", shard_index)
-
     @staticmethod
     def get_default_kwargs_eval(config: Config) -> Dict[str, Any]:
         """
@@ -118,8 +112,8 @@ def __init__(
         min_chunk_size=0,
         chunking_variance=0,
         estimated_num_seqs=None,
-        num_shards: int = 1,
-        shard_index: int = 0,
+        num_shards: Optional[int] = None,
+        shard_index: Optional[int] = None,
     ):
         """
         :param str name: e.g. "train" or "eval"
@@ -178,9 +172,9 @@ def __init__(
         self._chunking = chunking
         self.chunk_size, self.chunk_step, self.custom_chunking_func = self._parse_chunking(chunking)
         self._context_window = context_window
-        assert 0 <= shard_index < num_shards
-        self.num_shards = num_shards
-        self.shard_index = shard_index
+        assert (shard_index is None and num_shards is None) or 0 <= shard_index < num_shards
+        self._num_shards = num_shards
+        self._shard_index = shard_index
         if isinstance(context_window, (tuple, list)):
             assert len(context_window) == 2
             for elem in context_window:
@@ -219,7 +213,7 @@ def __repr__(self):
             getattr(self, "epoch", "<unknown>"),
         )
 
-    _getnewargs_exclude_attrs = set()  # type: typing.Set[str]
+    _getnewargs_exclude_attrs = {"num_shards", "shard_index"}  # type: typing.Set[str]
     _getnewargs_remap = {}  # type: typing.Dict[str,str]
 
     @staticmethod
@@ -256,25 +250,51 @@ def __reduce__(self):
         state = {attr: getattr(self, attr) for attr in ["epoch", "zpad"]}
         return Dataset._create_from_reduce, (self.__class__, kwargs, state)
 
+    @property
+    def num_shards(self) -> int:
+        if self._num_shards is None:
+            self._shard_index, self._num_shards = self._get_sharding_info()
+        return self._num_shards
+
+    @property
+    def shard_index(self) -> int:
+        if self._shard_index is None:
+            self._shard_index, self._num_shards = self._get_sharding_info()
+        return self._shard_index
+
     @staticmethod
-    def _get_rank_and_size(config: Config) -> Tuple[int, int]:
+    def _get_sharding_info(config: Optional[Config] = None) -> Tuple[int, int]:
         """
+        :param config: current RETURNN config, if not set, will fetch global
         :return: tuple (rank, size): the global rank and size for distributed trainings
         """
-        if config.typed_value("torch_distributed") is not None:
-            import returnn.torch.distributed
+        if config is None:
+            from returnn.config import get_global_config
 
-            ctx = returnn.torch.distributed.get_ctx(config=config)
-            return ctx.rank(), ctx.size()
-        elif config.is_true("use_horovod"):
+            config = get_global_config(return_empty_if_none=True)
+
+        if config.is_true("use_horovod"):
             assert config.bool("use_tensorflow", False) or config.value("backend", "").startswith("tensorflow")
 
             import returnn.tf.horovod
 
             ctx = returnn.tf.horovod.get_ctx(config=config)
             return ctx.rank(), ctx.size()
-        else:
-            return 0, 1
+
+        rank, size = 0, 1
+        if config.typed_value("torch_distributed") is not None:
+            import returnn.torch.distributed
+
+            ctx = returnn.torch.distributed.get_ctx(config=config)
+            rank, size = ctx.rank(), ctx.size()
+        if config.typed_value("torch_dataloader_opts") is not None:
+            import torch.utils.data
+
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                size *= worker_info.num_workers
+                rank += worker_info.id
+        return rank, size
 
     @property
     def random_seed_offset(self) -> int:

tests/test_torch_dataset.py

@@ -202,6 +202,20 @@ def test_MultiProcDataset_HDFDataset():
         assert c == n
 
 
+def test_dataset_num_workers_sharding():
+    config = Config({"backend": "torch", "torch_dataloader_opts": {"num_workers": 2}})
+    with global_config_ctx(config):
+        datasets = [
+            init_dataset({"class": "Task12AXDataset", "num_seqs": 100, "num_shards": 2, "shard_index": i})
+            for i in range(2)
+        ]
+        for dataset in datasets:
+            assert isinstance(dataset, Task12AXDataset)
+            dataset.init_seq_order(epoch=1)
+            assert dataset.shard_index < dataset.num_shards == 4
+            assert dataset.num_seqs == 25
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1: