Shuffled flatmap (meta-pytorch#1062)

PierreGtch · facebook-github-bot · commit 0061177c2e45 · 2023-03-17T18:21:52.000-07:00
Summary: Closes meta-pytorch#1043 ### Changes - Adds new datapipe `ShuffledFlatMapper`/`shuffled_flatmap` Pull Request resolved: meta-pytorch#1062 Reviewed By: NivekT Differential Revision: D44176076 Pulled By: ejguan fbshipit-source-id: a75f88fa3d28591d59f758d5affd6b1909841564
diff --git a/docs/source/torchdata.datapipes.iter.rst b/docs/source/torchdata.datapipes.iter.rst
@@ -167,6 +167,7 @@ These DataPipes apply the a given function to each element in the DataPipe.
     BatchMapper
     FlatMapper
     Mapper
+    ShuffledFlatMapper
     ThreadPoolMapper
 
 Other DataPipes
diff --git a/test/test_iterdatapipe.py b/test/test_iterdatapipe.py
@@ -894,6 +894,142 @@ def mul_fn(a, b):
         with self.assertRaisesRegex(TypeError, "length relies on the output of its function."):
             len(flatmapped_dp)
 
+    def test_shuffled_flatmap_iterdatapipe(self):
+        source_dp = IterableWrapper(list(range(20)))
+
+        def fn(e):
+            return [e, e * 10]
+
+        # Tests with buffer_size=1
+        # In this case, the expected behavior is similar to flatmap
+
+        shuffled_flatmapped_dp = source_dp.shuffled_flatmap(fn, buffer_size=1)
+        expected_list = list(itertools.chain(*[(e, e * 10) for e in source_dp]))
+
+        self.assertEqual(expected_list, list(shuffled_flatmapped_dp))
+
+        # Funtional Test: Specify input_col
+        tuple_source_dp = IterableWrapper([(d - 1, d, d + 1) for d in range(20)])
+
+        # Single input_col
+        input_col_1_dp = tuple_source_dp.shuffled_flatmap(fn, input_col=1, buffer_size=1)
+        self.assertEqual(expected_list, list(input_col_1_dp))
+
+        # With generator as fn
+        def gen_fn(e):
+            yield e
+            yield e * 10
+
+        shuffled_flatmapped_dp = source_dp.shuffled_flatmap(gen_fn, buffer_size=1)
+        expected_list = list(itertools.chain(*[(e, e * 10) for e in source_dp]))
+
+        self.assertEqual(expected_list, list(shuffled_flatmapped_dp))
+
+        # Multiple input_col
+        def mul_fn(a, b):
+            return [a - b, b - a]
+
+        input_col_2_dp = tuple_source_dp.shuffled_flatmap(mul_fn, input_col=(0, 2), buffer_size=1)
+        self.assertEqual(list(itertools.chain(*[(-2, 2) for _ in range(20)])), list(input_col_2_dp))
+
+        # shuffled_flatmap with no fn specified
+        default_dp = tuple_source_dp.shuffled_flatmap(buffer_size=1)
+        self.assertEqual(list(itertools.chain(*[(n - 1, n, n + 1) for n in range(20)])), list(default_dp))
+
+        # shuffled_flatmap with no fn specified, multiple input_col
+        default_dp = tuple_source_dp.shuffled_flatmap(input_col=(0, 2), buffer_size=1)
+        self.assertEqual(list(itertools.chain(*[(n - 1, n + 1) for n in range(20)])), list(default_dp))
+
+        # shuffled_flatmap with no fn specified, some special input
+        tuple_source_dp = IterableWrapper([[1, 2, [3, 4]], [5, 6, [7, 8]]])
+        default_dp = tuple_source_dp.shuffled_flatmap(input_col=(0, 2), buffer_size=1)
+        self.assertEqual([1, [3, 4], 5, [7, 8]], list(default_dp))
+
+        # Reset Test: reset the DataPipe after reading part of it
+        n_elements_before_reset = 5
+        res_before_reset, res_after_reset = reset_after_n_next_calls(shuffled_flatmapped_dp, n_elements_before_reset)
+
+        self.assertEqual(expected_list[:n_elements_before_reset], res_before_reset)
+        self.assertEqual(expected_list, res_after_reset)
+
+        # __len__ Test: length should be len(source_dp)*len(fn->out_shape) which we can't know
+        with self.assertRaisesRegex(TypeError, "length relies on the output of its function."):
+            len(shuffled_flatmapped_dp)
+
+        # __len__ when no fn specified:
+        dp = IterableWrapper([[1, 2], [], [3], [4, 5, 6, [7, 8]]])
+        dp = dp.shuffled_flatmap()
+        self.assertEqual(len(dp), 7)
+
+        # Tests with .set_shuffle(False)
+        # In this case, the expected behavior is similar to flatmap
+
+        shuffled_flatmapped_dp = source_dp.shuffled_flatmap(fn).set_shuffle(False)
+        expected_list = list(itertools.chain(*[(e, e * 10) for e in source_dp]))
+
+        self.assertEqual(expected_list, list(shuffled_flatmapped_dp))
+
+        # Funtional Test: Specify input_col
+        tuple_source_dp = IterableWrapper([(d - 1, d, d + 1) for d in range(20)])
+
+        # Single input_col
+        input_col_1_dp = tuple_source_dp.shuffled_flatmap(fn, input_col=1, buffer_size=1)
+        self.assertEqual(expected_list, list(input_col_1_dp))
+
+        # Multiple input_col
+        input_col_2_dp = tuple_source_dp.shuffled_flatmap(mul_fn, input_col=(0, 2)).set_shuffle(False)
+        self.assertEqual(list(itertools.chain(*[(-2, 2) for _ in range(20)])), list(input_col_2_dp))
+
+        # shuffled_flatmap with no fn specified
+        default_dp = tuple_source_dp.shuffled_flatmap().set_shuffle(False)
+        self.assertEqual(list(itertools.chain(*[(n - 1, n, n + 1) for n in range(20)])), list(default_dp))
+
+        # shuffled_flatmap with no fn specified, multiple input_col
+        default_dp = tuple_source_dp.shuffled_flatmap(input_col=(0, 2)).set_shuffle(False)
+        self.assertEqual(list(itertools.chain(*[(n - 1, n + 1) for n in range(20)])), list(default_dp))
+
+        # shuffled_flatmap with no fn specified, some special input
+        tuple_source_dp = IterableWrapper([[1, 2, [3, 4]], [5, 6, [7, 8]]])
+        default_dp = tuple_source_dp.shuffled_flatmap(input_col=(0, 2)).set_shuffle(False)
+        self.assertEqual([1, [3, 4], 5, [7, 8]], list(default_dp))
+
+        # Reset Test: reset the DataPipe after reading part of it
+        n_elements_before_reset = 5
+        res_before_reset, res_after_reset = reset_after_n_next_calls(shuffled_flatmapped_dp, n_elements_before_reset)
+
+        self.assertEqual(expected_list[:n_elements_before_reset], res_before_reset)
+        self.assertEqual(expected_list, res_after_reset)
+
+        # Other tests
+
+        # Test no empty buffers:
+        with self.assertRaises(AssertionError):
+            _ = source_dp.shuffled_flatmap(buffer_size=0)
+
+        # Functional Test: No seed
+        consecutive_tuple_source_dp = IterableWrapper([(d, d + 1, d + 2) for d in range(0, 21, 3)])
+        shuffled_flatmapped_dp = consecutive_tuple_source_dp.shuffled_flatmap()
+        self.assertEqual(set(range(21)), set(shuffled_flatmapped_dp))
+
+        # Functional Test: With global seed
+        torch.manual_seed(123)
+        shuffled_flatmapped_dp = tuple_source_dp.shuffled_flatmap()
+        res = list(shuffled_flatmapped_dp)
+        torch.manual_seed(123)
+        self.assertEqual(list(shuffled_flatmapped_dp), res)
+
+        # Functional Test: Set seed
+        shuffled_flatmapped_dp = tuple_source_dp.shuffled_flatmap().set_seed(123)
+        res = list(shuffled_flatmapped_dp)
+        shuffled_flatmapped_dp.set_seed(123)
+        self.assertEqual(list(shuffled_flatmapped_dp), res)
+
+        # Reset Test:
+        shuffled_flatmapped_dp = tuple_source_dp.shuffled_flatmap()
+        n_elements_before_reset = 5
+        res_before_reset, res_after_reset = reset_after_n_next_calls(shuffled_flatmapped_dp, n_elements_before_reset)
+        self.assertEqual(5, len(res_before_reset))
+
     def test_round_robin_demux_iterdatapipe(self):
         source_dp = IterableWrapper(list(range(23)))
         with self.assertRaisesRegex(ValueError, "Expected `num_instaces`"):
diff --git a/test/test_serialization.py b/test/test_serialization.py
@@ -197,6 +197,7 @@ def test_serializable(self):
             (iterdp.Dropper, IterableWrapper([(0, 0), (0, 0), (0, 0), (0, 0)]), ([1]), {}),
             (iterdp.Enumerator, None, (2,), {}),
             (iterdp.FlatMapper, None, (_fake_fn_ls,), {}),
+            (iterdp.ShuffledFlatMapper, None, (_fake_fn_ls,), {"buffer_size": 1}),
             (iterdp.Flattener, IterableWrapper([(0, (0, 1)), (0, (0, 1)), (0, (0, 1)), (0, (0, 1))]), ([1]), {}),
             (iterdp.FSSpecFileLister, ".", (), {}),
             (iterdp.FSSpecFileOpener, None, (), {}),
@@ -363,6 +364,7 @@ def test_serializable_with_dill(self):
         unpicklable_datapipes: List = [
             (iterdp.BatchMapper, (lambda batch: [d + 1 for d in batch], 2), {}),
             (iterdp.FlatMapper, (lambda x: [x, x],), {}),
+            (iterdp.ShuffledFlatMapper, (lambda x: [x, x],), {"buffer_size": 1}),
             (iterdp.IterKeyZipper, (ref_idp, lambda x: x, None, True, 100), {}),
             (iterdp.MapKeyZipper, (ref_mdp, lambda x: x), {}),
             (iterdp.OnDiskCacheHolder, (lambda x: x,), {}),
diff --git a/torchdata/datapipes/iter/__init__.py b/torchdata/datapipes/iter/__init__.py
@@ -71,6 +71,7 @@
     DropperIterDataPipe as Dropper,
     FlatMapperIterDataPipe as FlatMapper,
     FlattenIterDataPipe as Flattener,
+    ShuffledFlatMapperIterDataPipe as ShuffledFlatMapper,
     SliceIterDataPipe as Slicer,
     ThreadPoolMapperIterDataPipe as ThreadPoolMapper,
 )
@@ -209,6 +210,7 @@
     "ShardExpander",
     "ShardingFilter",
     "ShardingRoundRobinDispatcher",
+    "ShuffledFlatMapper",
     "Shuffler",
     "Slicer",
     "StreamReader",
diff --git a/torchdata/datapipes/iter/transform/callable.py b/torchdata/datapipes/iter/transform/callable.py
@@ -6,12 +6,14 @@
 
 import asyncio
 import inspect
+import random
 import warnings
 from collections import deque
 from concurrent import futures
 
 from typing import Callable, Hashable, Iterator, List, Optional, Set, Sized, TypeVar, Union
 
+import torch
 from torch.utils.data.datapipes.utils.common import _check_unpickable_fn, validate_input_col
 from torchdata.datapipes import functional_datapipe
 from torchdata.datapipes.iter import IterDataPipe
@@ -170,6 +172,161 @@ def __len__(self) -> int:
         raise TypeError(f"{type(self).__name__}'s length relies on the output of its function.")
 
 
+@functional_datapipe("shuffled_flatmap")
+class ShuffledFlatMapperIterDataPipe(IterDataPipe):
+    r"""
+    Applies a function over each item from the source DataPipe,
+    then collects the iterables returned in a buffer,
+    then, at every iteration, chooses at random one of the iterables in the buffer
+    and yields one item from this iterable (functional name: ``shuffled_flatmap``).
+
+    When the buffer is full, the DataPipe will begin to yield elements from iterables within the buffer.
+    New iterables will be added to the buffer once the existing ones run out of elements.
+    Note:
+        The output from ``fn`` must be an Iterable. Otherwise, an error will be raised.
+        If ``fn`` is ``None``, source DataPipe will be just flattened vertically, provided that items can be unpacked.
+
+    Args:
+        datapipe: Source IterDataPipe
+        fn: the function to be applied to each element in the DataPipe, the output must be a Sequence
+        input_col: Index or indices of data which ``fn`` is applied, such as:
+
+            - ``None`` as default to apply ``fn`` to the data directly.
+            - Integer(s) is/are used for list/tuple.
+            - Key(s) is/are used for dict.
+        buffer_size: the max number of iterables this DataPipe can hold at a time (default to ``100``)
+
+    Example:
+        >>> from torchdata.datapipes.iter import IterableWrapper
+        >>> source_dp = IterableWrapper([[1, 2, 3, 4], 'abcd', 'ABCD'])
+        >>> shuffled_flatmapped_dp = source_dp.shuffled_flatmap(buffer_size=2)
+        >>> list(shuffled_flatmapped_dp)
+        ['a', 'b', 'c', 1, 'd', 'A', 'B', 'C', 2, 'D', 3, 4]
+        >>>
+        >>> # To shuffle all the elements, you can combine `shuffled_flatmap` with `in_batch_shuffle` like this:
+        >>> fully_shuffled_flatmapped_dp = source_dp.in_batch_shuffle()
+        >>> fully_shuffled_flatmapped_dp = fully_shuffled_flatmapped_dp.shuffled_flatmap()
+        >>> list(fully_shuffled_flatmapped_dp)
+        ['b', 3, 'c', 'd', 'C', 'A', 'a', 2, 'B', 'D', 4, 1]
+    """
+    datapipe: IterDataPipe
+    fn: Optional[Callable]
+    buffer_size: int
+    _buffer: List[Iterator]
+    _enabled: bool
+    _seed: Optional[int]
+    _rng: random.Random
+    _no_op_fn: bool = False
+
+    def __init__(
+        self, datapipe: IterDataPipe, fn: Optional[Callable] = None, input_col=None, buffer_size: int = 100
+    ) -> None:
+        super().__init__()
+        self._buffer = []
+        self.datapipe = datapipe
+
+        if fn is None:
+            fn = _no_op_fn
+            self._no_op_fn = True
+        _check_unpickable_fn(fn)
+        self.fn = fn  # type: ignore[assignment]
+        self.input_col = input_col
+        validate_input_col(fn, input_col)
+
+        assert buffer_size > 0, "buffer_size should be larger than 0"
+        self.buffer_size = buffer_size
+        self._enabled = True
+        self._seed = None
+        self._rng = random.Random()
+
+    def set_shuffle(self, shuffle=True):
+        self._enabled = shuffle
+        return self
+
+    def set_seed(self, seed: int):
+        self._seed = seed
+        return self
+
+    def reset(self) -> None:
+        self._buffer = []
+        if self._enabled:
+            if self._seed is None:
+                self._seed = int(torch.empty((), dtype=torch.int64).random_().item())
+            self._rng.seed(self._seed)
+            self._seed = None
+
+    def _apply_fn(self, data):
+        if self.input_col is None:
+            return self.fn(data)  # type: ignore[misc]
+        elif isinstance(self.input_col, (list, tuple)):
+            args = tuple(data[col] for col in self.input_col)
+            return self.fn(*args)  # type: ignore[misc]
+        else:
+            return self.fn(data[self.input_col])  # type: ignore[misc]
+
+    def __iter__(self) -> Iterator[T_co]:
+        if not self._enabled:  # equivalent to flatmap
+            for x in self.datapipe:
+                yield from self._apply_fn(x)
+        else:
+            idx = self._rng.randint(0, self.buffer_size - 1)
+            for x in self.datapipe:
+                while len(self._buffer) == self.buffer_size:
+                    try:
+                        yield next(self._buffer[idx])
+                        idx = self._rng.randint(0, self.buffer_size - 1)
+                    except StopIteration:
+                        self._buffer.pop(idx)
+                self._buffer.append(iter(self._apply_fn(x)))
+            while self._buffer:
+                try:
+                    idx = self._rng.randint(0, len(self._buffer) - 1)
+                    yield next(self._buffer[idx])
+                except StopIteration:
+                    self._buffer.pop(idx)
+
+    def __len__(self) -> int:
+        if self._no_op_fn:
+            return sum(map(len, self.datapipe))
+        raise TypeError(f"{type(self).__name__}'s length relies on the output of its function.")
+
+    def __getstate__(self):
+        state = (
+            self.datapipe,
+            self.fn,
+            self.input_col,
+            self.buffer_size,
+            self._buffer,
+            self._enabled,
+            self._seed,
+            self._rng.getstate(),
+            self._valid_iterator_id,
+            self._number_of_samples_yielded,
+        )
+        if IterDataPipe.getstate_hook is not None:
+            return IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        (
+            self.datapipe,
+            self.fn,
+            self.input_col,
+            self.buffer_size,
+            self._buffer,
+            self._enabled,
+            self._seed,
+            rng_state,
+            self._valid_iterator_id,
+            self._number_of_samples_yielded,
+        ) = state
+        self._rng = random.Random()
+        self._rng.setstate(rng_state)
+
+    def __del__(self):
+        self._buffer.clear()
+
+
 @functional_datapipe("drop")
 class DropperIterDataPipe(IterDataPipe[T_co]):
     r"""
