Repository: pytorch/ignite

Issue:
Make idist.auto_dataloader work on IterableDataset
## 🚀 Feature

Following https://github.com/pytorch/ignite/issues/618#issuecomment-852082567 we need to fix `idist.auto_dataloader` to make it work when dataset is IterableDataset. Basically, we probably have to remove the sampler.


Discussion:

Resolving patch:
diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py
--- a/ignite/distributed/auto.py
+++ b/ignite/distributed/auto.py
@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 from torch.optim.optimizer import Optimizer
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data.sampler import Sampler
 
@@ -50,7 +50,9 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> Union[DataLoader, "_MpDe
         )
 
     Args:
-        dataset: input torch dataset
+        dataset: input torch dataset. If input dataset is `torch IterableDataset`_ then dataloader will be
+            created without any distributed sampling. Please, make sure that the dataset itself produces
+            different data on different ranks.
         kwargs: keyword arguments for `torch DataLoader`_.
 
     Returns:
@@ -60,6 +62,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> Union[DataLoader, "_MpDe
     .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178
     .. _torch DistributedSampler:
         https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
+    .. _torch IterableDataset: https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset
     """
     rank = idist.get_rank()
     world_size = idist.get_world_size()
@@ -74,19 +77,25 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> Union[DataLoader, "_MpDe
             kwargs["num_workers"] = (kwargs["num_workers"] + nproc - 1) // nproc
 
         if "batch_sampler" not in kwargs:
-            if kwargs.get("sampler", None) is not None:
-                sampler = DistributedProxySampler(
-                    kwargs["sampler"], num_replicas=world_size, rank=rank
-                )  # type: Union[DistributedProxySampler, DistributedSampler, Sampler]
-            else:
-                sampler = DistributedSampler(
-                    dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get("shuffle", True)
+            if isinstance(dataset, IterableDataset):
+                logger.info(
+                    "Found iterable dataset, dataloader will be created without any distributed sampling. "
+                    "Please, make sure that the dataset itself produces different data on different ranks."
                 )
-                # we need to remove "shuffle" from kwargs if sampler is used
-                if "shuffle" in kwargs:
-                    del kwargs["shuffle"]
-
-            kwargs["sampler"] = sampler
+            else:
+                if kwargs.get("sampler", None) is not None:
+                    sampler = DistributedProxySampler(
+                        kwargs["sampler"], num_replicas=world_size, rank=rank
+                    )  # type: Union[DistributedProxySampler, DistributedSampler, Sampler]
+                else:
+                    sampler = DistributedSampler(
+                        dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get("shuffle", True)
+                    )
+                    # we need to remove "shuffle" from kwargs if sampler is used
+                    if "shuffle" in kwargs:
+                        del kwargs["shuffle"]
+
+                kwargs["sampler"] = sampler
         else:
             warnings.warn(
                 "Found batch_sampler in provided kwargs. Please, make sure that it is compatible "


Test patch:
diff --git a/tests/ignite/distributed/test_auto.py b/tests/ignite/distributed/test_auto.py
--- a/tests/ignite/distributed/test_auto.py
+++ b/tests/ignite/distributed/test_auto.py
@@ -5,7 +5,8 @@
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader
-from torch.utils.data.dataset import Dataset
+from torch.utils.data.dataloader import _InfiniteConstantSampler
+from torch.utils.data.dataset import Dataset, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data.sampler import BatchSampler, RandomSampler, Sampler, SequentialSampler, WeightedRandomSampler
 
@@ -18,6 +19,19 @@ def __getitem__(self, index):
         return index
 
 
+class DummyIterableDataset(IterableDataset):
+    def __init__(self, start, end):
+        super(DummyIterableDataset).__init__()
+        self.start = start
+        self.end = end
+
+    def __iter__(self):
+        return iter(range(self.start, self.end))
+
+    def __len__(self):
+        return self.end - self.start
+
+
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif("WORLD_SIZE" not in os.environ, reason="Skip if WORLD_SIZE not in env vars")
@@ -37,42 +51,51 @@ def test_auto_dataloader_warning_tpu():
 
 
 def _test_auto_dataloader(ws, nproc, batch_size, num_workers=1, sampler_name=None, dl_type=DataLoader):
+    def _test(data):
+        if sampler_name is None:
+            sampler = None
+        elif sampler_name == "WeightedRandomSampler":
+            sampler = WeightedRandomSampler(weights=torch.ones(100), num_samples=100)
+        else:
+            raise RuntimeError(f"Unknown sampler name: {sampler_name}")
+
+        # Test auto_dataloader
+        assert idist.get_world_size() == ws, f"{idist.get_world_size()} vs {ws}"
+
+        shuffle = sampler is None if not isinstance(data, IterableDataset) else False
+        dataloader = auto_dataloader(
+            data, batch_size=batch_size, num_workers=num_workers, sampler=sampler, shuffle=shuffle
+        )
 
-    data = torch.rand(100, 3, 12, 12)
+        assert isinstance(dataloader, dl_type)
+        if hasattr(dataloader, "_loader"):
+            dataloader = dataloader._loader
+        if ws < batch_size:
+            assert dataloader.batch_size == batch_size // ws
+        else:
+            assert dataloader.batch_size == batch_size
+        if ws <= num_workers:
+            assert dataloader.num_workers == (num_workers + nproc - 1) // nproc
+        else:
+            assert dataloader.num_workers == num_workers
+
+        if isinstance(data, IterableDataset):
+            sampler_type = _InfiniteConstantSampler
+        elif ws > 1:
+            sampler_type = DistributedSampler if sampler is None else DistributedProxySampler
+        else:
+            sampler_type = RandomSampler if sampler is None else type(sampler)
 
-    if sampler_name is None:
-        sampler = None
-    elif sampler_name == "WeightedRandomSampler":
-        sampler = WeightedRandomSampler(weights=torch.ones(100), num_samples=100)
-    else:
-        raise RuntimeError(f"Unknown sampler name: {sampler_name}")
-
-    # Test auto_dataloader
-    assert idist.get_world_size() == ws, f"{idist.get_world_size()} vs {ws}"
-    dataloader = auto_dataloader(
-        data, batch_size=batch_size, num_workers=num_workers, sampler=sampler, shuffle=sampler is None
-    )
-
-    assert isinstance(dataloader, dl_type)
-    if hasattr(dataloader, "_loader"):
-        dataloader = dataloader._loader
-    if ws < batch_size:
-        assert dataloader.batch_size == batch_size // ws
-    else:
-        assert dataloader.batch_size == batch_size
-    if ws <= num_workers:
-        assert dataloader.num_workers == (num_workers + nproc - 1) // nproc
-    else:
-        assert dataloader.num_workers == num_workers
-
-    if ws < 2:
-        sampler_type = RandomSampler if sampler is None else type(sampler)
-        assert isinstance(dataloader.sampler, sampler_type)
-    else:
-        sampler_type = DistributedSampler if sampler is None else DistributedProxySampler
         assert isinstance(dataloader.sampler, sampler_type)
-    if isinstance(dataloader, DataLoader):
-        assert dataloader.pin_memory == ("cuda" in idist.device().type)
+
+        if isinstance(dataloader, DataLoader):
+            assert dataloader.pin_memory == ("cuda" in idist.device().type)
+
+    data = torch.rand(100, 3, 12, 12)
+    _test(data)
+    if sampler_name is None:
+        data = DummyIterableDataset(0, 100)
+        _test(data)
 
 
 def _test_auto_model(model, ws, device, sync_bn=False, **kwargs):