Repository: Project-MONAI/MONAI

Issue:
Support of torch.utils.data.IterableDataset (and WebDataset?)
**Is your feature request related to a problem? Please describe.**
`monai.data.Dataset` mainly supports `torch.utils.data.Dataset` with the `__len__` and `__getitem__` APIs.

would be great to introduce compatible APIs based on `torch.utils.data.IterableDataset`,
the intention is to support the WebDataset https://github.com/pytorch/pytorch/issues/38419 in the future, for efficient I/O of very large datasets

see also https://github.com/tmbdev/webdataset


Discussion:

Resolving patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -24,6 +24,7 @@
 from .decathlon_datalist import load_decathlon_datalist, load_decathlon_properties
 from .grid_dataset import *
 from .image_reader import *
+from .iterable_dataset import IterableDataset
 from .nifti_reader import NiftiDataset
 from .nifti_saver import NiftiSaver
 from .nifti_writer import write_nifti
diff --git a/monai/data/iterable_dataset.py b/monai/data/iterable_dataset.py
new file mode 100644
--- /dev/null
+++ b/monai/data/iterable_dataset.py
@@ -0,0 +1,46 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Iterable, Optional
+
+from torch.utils.data import IterableDataset as _TorchIterableDataset
+
+from monai.transforms import apply_transform
+
+
+class IterableDataset(_TorchIterableDataset):
+    """
+    A generic dataset for iterable data source and an optional callable data transform
+    when fetching a data sample.
+    For example, typical input data can be web data stream which can support multi-process access.
+
+    Note that when used with `DataLoader` and `num_workers > 0`, each worker process will have a
+    different copy of the dataset object, need to guarantee process-safe from data source or DataLoader.
+
+    """
+
+    def __init__(self, data: Iterable, transform: Optional[Callable] = None) -> None:
+        """
+        Args:
+            data: input data source to load and transform to generate dataset for model.
+            transform: a callable data transform on input data.
+        """
+        self.data = data
+        self.transform = transform
+        self.source = None
+
+    def __iter__(self):
+        self.source = iter(self.data)
+        for data in self.source:
+            if self.transform is not None:
+                data = apply_transform(self.transform, data)
+
+            yield data


Test patch:
diff --git a/tests/min_tests.py b/tests/min_tests.py
--- a/tests/min_tests.py
+++ b/tests/min_tests.py
@@ -32,6 +32,7 @@ def run_testsuit():
         "test_cachedataset",
         "test_cachedataset_parallel",
         "test_dataset",
+        "test_iterable_dataset",
         "test_ensemble_evaluator",
         "test_handler_checkpoint_loader",
         "test_handler_checkpoint_saver",
diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py
new file mode 100644
--- /dev/null
+++ b/tests/test_iterable_dataset.py
@@ -0,0 +1,88 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import tempfile
+import unittest
+from multiprocessing import Lock
+
+import nibabel as nib
+import numpy as np
+
+from monai.data import DataLoader, IterableDataset
+from monai.transforms import Compose, LoadNiftid, SimulateDelayd
+
+lock = Lock()
+
+
+class _Stream:
+    def __init__(self, data, dbpath):
+        # simulate a database at the website
+        self.dbpath = dbpath
+        self.reset()
+        self.data = data
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        data = None
+        # support multi-process access to the database
+        lock.acquire()
+        with open(self.dbpath) as f:
+            count = json.load(f)["count"]
+            if count > 0:
+                data = self.data[count - 1]
+                with open(self.dbpath, "w") as f:
+                    json.dump({"count": count - 1}, f)
+        lock.release()
+
+        if count == 0:
+            raise StopIteration
+        return data
+
+    def reset(self):
+        with open(self.dbpath, "w") as f:
+            json.dump({"count": 6}, f)
+
+
+class TestIterableDataset(unittest.TestCase):
+    def test_shape(self):
+        expected_shape = (128, 128, 128)
+        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
+        test_data = list()
+        with tempfile.TemporaryDirectory() as tempdir:
+            for i in range(6):
+                nib.save(test_image, os.path.join(tempdir, f"test_image{str(i)}.nii.gz"))
+                test_data.append({"image": os.path.join(tempdir, f"test_image{str(i)}.nii.gz")})
+
+            test_transform = Compose(
+                [
+                    LoadNiftid(keys="image"),
+                    SimulateDelayd(keys="image", delay_time=1e-7),
+                ]
+            )
+
+            test_stream = _Stream(data=test_data, dbpath=os.path.join(tempdir, "countDB"))
+
+            dataset = IterableDataset(data=test_stream, transform=test_transform)
+            for d in dataset:
+                self.assertTupleEqual(d["image"].shape, expected_shape)
+
+            test_stream.reset()
+            dataloader = DataLoader(dataset=dataset, batch_size=3, num_workers=2)
+            for d in dataloader:
+                self.assertTupleEqual(d["image"].shape[1:], expected_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()