From 0f2c60f0e37fe1738d30d3db3fd48e88b8087af7 Mon Sep 17 00:00:00 2001 From: Vitaly Fedyunin Date: Fri, 20 Aug 2021 08:36:14 -0700 Subject: [PATCH] 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 --- torch/utils/data/datapipes/iter/__init__.py | 4 ++++ torch/utils/data/datapipes/iter/utils.py | 10 ++++++++++ 2 files changed, 14 insertions(+) create mode 100644 torch/utils/data/datapipes/iter/utils.py diff --git a/torch/utils/data/datapipes/iter/__init__.py b/torch/utils/data/datapipes/iter/__init__.py index 0bcfdc4..bdaef95 100644 --- 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 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 -- 2.7.4