Args: size (int): the total number of data of the underlying dataset to sample from shuffle (bool): whether to shuffle the indices or not seed (int): the initial seed of the shuffle. Must be the same across all workers. If None, will use a
(
self,
size: int,
shuffle: bool = True,
seed: Optional[int] = 0,
rank=0,
world_size=1,
)
| 49 | """ |
| 50 | |
| 51 | def __init__( |
| 52 | self, |
| 53 | size: int, |
| 54 | shuffle: bool = True, |
| 55 | seed: Optional[int] = 0, |
| 56 | rank=0, |
| 57 | world_size=1, |
| 58 | ): |
| 59 | """ |
| 60 | Args: |
| 61 | size (int): the total number of data of the underlying dataset to sample from |
| 62 | shuffle (bool): whether to shuffle the indices or not |
| 63 | seed (int): the initial seed of the shuffle. Must be the same |
| 64 | across all workers. If None, will use a random seed shared |
| 65 | among workers (require synchronization among all workers). |
| 66 | """ |
| 67 | self._size = size |
| 68 | assert size > 0 |
| 69 | self._shuffle = shuffle |
| 70 | self._seed = int(seed) |
| 71 | |
| 72 | if dist.is_available() and dist.is_initialized(): |
| 73 | self._rank = dist.get_rank() |
| 74 | self._world_size = dist.get_world_size() |
| 75 | else: |
| 76 | self._rank = rank |
| 77 | self._world_size = world_size |
| 78 | |
| 79 | def __iter__(self): |
| 80 | start = self._rank |