Adding IterableAsDataPipe IterDataPipe (#63522)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63522

Supports sharding and batching on loader level
* **#63522 Adding IterableAsDataPipe IterDataPipe
usefull for tests and simple cases**

usefull for tests and simple cases

Test Plan: Imported from OSS

Reviewed By: ejguan

Differential Revision: D30426528

Pulled By: VitalyFedyunin

fbshipit-source-id: 535b5cc1505bb58731fcca8170541ac5ee7bd417

diff --git a/torch/utils/data/datapipes/iter/__init__.py b/torch/utils/data/datapipes/iter/__init__.py
index 0bcfdc4..bdaef95 100644 (file)
--- a/torch/utils/data/datapipes/iter/__init__.py
+++ b/torch/utils/data/datapipes/iter/__init__.py
@@ -43,6 +43,9 @@ from torch.utils.data.datapipes.iter.selecting import (
 from torch.utils.data.datapipes.iter.tobytes import (
     ToBytesIterDataPipe as ToBytes,
 )
+from torch.utils.data.datapipes.iter.utils import (
+    IterableAsDataPipeIterDataPipe as IterableAsDataPipe,
+)
 
 __all__ = ['Batch',
            'BucketBatcher',
@@ -51,6 +54,7 @@ __all__ = ['Batch',
            'Filter',
            'GroupByKey',
            'HttpReader',
+           'IterableAsDataPipe',
            'ListDirFiles',
            'LoadFilesFromDisk',
            'Map',

diff --git a/torch/utils/data/datapipes/iter/utils.py b/torch/utils/data/datapipes/iter/utils.py
new file mode 100644 (file)
index 0000000..ea241d9
--- /dev/null
+++ b/torch/utils/data/datapipes/iter/utils.py
@@ -0,0 +1,10 @@
+from torch.utils.data import IterDataPipe
+
+
+class IterableAsDataPipeIterDataPipe(IterDataPipe):
+    def __init__(self, iterable):
+        self.iterable = iterable
+
+    def __iter__(self):
+        for data in self.iterable:
+            yield data