1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
# mypy: allow-untyped-defs
import random
from typing import Iterator, List, Optional, TypeVar
import torch
from torch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe
__all__ = ["ShufflerIterDataPipe"]
_T_co = TypeVar("_T_co", covariant=True)
# @functional_datapipe('shuffle')
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)
|
