| 81 | """ |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | dataset: Dataset, |
| 86 | shuffle: bool = True, |
| 87 | seed: int = 0, |
| 88 | drop_last: bool = False, |
| 89 | ) -> None: |
| 90 | self.dataset = dataset |
| 91 | self.epoch = 0 |
| 92 | self.idx = 0 |
| 93 | self.drop_last = drop_last |
| 94 | self.world_size = dist.get_world_size() if dist.is_initialized() else 1 |
| 95 | self.rank = dist.get_rank() if dist.is_initialized() else 0 |
| 96 | # If the dataset length is evenly divisible by # of replicas, then there |
| 97 | # is no need to drop any data, since the dataset will be split equally. |
| 98 | if self.drop_last and len(self.dataset) % self.world_size != 0: # type: ignore[arg-type] |
| 99 | # Split to nearest available length that is evenly divisible. |
| 100 | # This is to ensure each rank receives the same amount of data when |
| 101 | # using this Sampler. |
| 102 | self.num_samples = math.ceil( |
| 103 | (len(self.dataset) - self.world_size) / self.world_size # type: ignore[arg-type] |
| 104 | ) |
| 105 | else: |
| 106 | self.num_samples = math.ceil(len(self.dataset) / self.world_size) # type: ignore[arg-type] |
| 107 | self.total_size = self.num_samples * self.world_size |
| 108 | self.shuffle = shuffle |
| 109 | self.seed = seed |
| 110 | |
| 111 | def __iter__(self) -> Iterator: |
| 112 | if self.shuffle: |