|
|
import random |
|
|
|
|
|
import torch |
|
|
from torch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe |
|
|
from typing import Iterator, List, Optional, TypeVar |
|
|
|
|
|
__all__ = ["ShufflerIterDataPipe", ] |
|
|
|
|
|
|
|
|
T_co = TypeVar('T_co', covariant=True) |
|
|
|
|
|
|
|
|
|
|
|
class ShufflerIterDataPipe(IterDataPipe[T_co]): |
|
|
r""" |
|
|
Shuffle the input MapDataPipe via its indices (functional name: ``shuffle``). |
|
|
|
|
|
When it is used with :class:`~torch.utils.data.DataLoader`, the methods to |
|
|
set up random seed are different based on :attr:`num_workers`. |
|
|
|
|
|
For single-process mode (:attr:`num_workers == 0`), the random seed is set before |
|
|
the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process |
|
|
mode (:attr:`num_worker > 0`), ``worker_init_fn`` is used to set up a random seed |
|
|
for each worker process. |
|
|
|
|
|
Args: |
|
|
datapipe: MapDataPipe being shuffled |
|
|
indices: a list of indices of the MapDataPipe. If not provided, we assume it uses 0-based indexing |
|
|
|
|
|
Example: |
|
|
>>> # xdoctest: +SKIP |
|
|
>>> from torchdata.datapipes.map import SequenceWrapper |
|
|
>>> dp = SequenceWrapper(range(10)) |
|
|
>>> shuffle_dp = dp.shuffle().set_seed(0) |
|
|
>>> list(shuffle_dp) |
|
|
[7, 8, 1, 5, 3, 4, 2, 0, 9, 6] |
|
|
>>> list(shuffle_dp) |
|
|
[6, 1, 9, 5, 2, 4, 7, 3, 8, 0] |
|
|
>>> # Reset seed for Shuffler |
|
|
>>> shuffle_dp = shuffle_dp.set_seed(0) |
|
|
>>> list(shuffle_dp) |
|
|
[7, 8, 1, 5, 3, 4, 2, 0, 9, 6] |
|
|
|
|
|
Note: |
|
|
Even thought this ``shuffle`` operation takes a ``MapDataPipe`` as the input, it would return an |
|
|
``IterDataPipe`` rather than a ``MapDataPipe``, because ``MapDataPipe`` should be non-sensitive to |
|
|
the order of data order for the sake of random reads, but ``IterDataPipe`` depends on the order |
|
|
of data during data-processing. |
|
|
""" |
|
|
datapipe: MapDataPipe[T_co] |
|
|
_enabled: bool |
|
|
_seed: Optional[int] |
|
|
_rng: random.Random |
|
|
|
|
|
def __init__(self, |
|
|
datapipe: MapDataPipe[T_co], |
|
|
*, |
|
|
indices: Optional[List] = None, |
|
|
) -> None: |
|
|
super().__init__() |
|
|
self.datapipe = datapipe |
|
|
self.indices = list(range(len(datapipe))) if indices is None else indices |
|
|
self._enabled = True |
|
|
self._seed = None |
|
|
self._rng = random.Random() |
|
|
self._shuffled_indices: List = self.indices |
|
|
|
|
|
def set_shuffle(self, shuffle=True): |
|
|
self._enabled = shuffle |
|
|
return self |
|
|
|
|
|
def set_seed(self, seed: int): |
|
|
self._seed = seed |
|
|
return self |
|
|
|
|
|
def __iter__(self) -> Iterator[T_co]: |
|
|
if not self._enabled: |
|
|
for idx in self.indices: |
|
|
yield self.datapipe[idx] |
|
|
else: |
|
|
while self._shuffled_indices: |
|
|
idx = self._shuffled_indices.pop() |
|
|
yield self.datapipe[idx] |
|
|
|
|
|
def reset(self) -> None: |
|
|
if self._enabled and self._seed is None: |
|
|
self._seed = int(torch.empty((), dtype=torch.int64).random_().item()) |
|
|
self._rng.seed(self._seed) |
|
|
self._seed = None |
|
|
self._shuffled_indices = self._rng.sample(self.indices, len(self.indices)) |
|
|
|
|
|
def __len__(self) -> int: |
|
|
return len(self.datapipe) |
|
|
|
|
|
def __getstate__(self): |
|
|
state = ( |
|
|
self.datapipe, |
|
|
self.indices, |
|
|
self._enabled, |
|
|
self._seed, |
|
|
self._rng.getstate(), |
|
|
self._shuffled_indices, |
|
|
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.indices, |
|
|
self._enabled, |
|
|
self._seed, |
|
|
rng_state, |
|
|
self._shuffled_indices, |
|
|
self._valid_iterator_id, |
|
|
self._number_of_samples_yielded, |
|
|
) = state |
|
|
self._rng = random.Random() |
|
|
self._rng.setstate(rng_state) |
|
|
|
|
|
|
|
|
MapDataPipe.register_datapipe_as_function("shuffle", ShufflerIterDataPipe) |
|
|
|
