| 94 | self.seed = seed |
| 95 | |
| 96 | def __iter__(self) -> Iterator[T_co]: |
| 97 | if self.shuffle: |
| 98 | # deterministically shuffle based on epoch and seed |
| 99 | g = torch.Generator() |
| 100 | g.manual_seed(self.seed + self.epoch) |
| 101 | indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] |
| 102 | else: |
| 103 | indices = list(range(len(self.dataset))) # type: ignore[arg-type] |
| 104 | |
| 105 | if not self.drop_last: |
| 106 | # add extra samples to make it evenly divisible |
| 107 | padding_size = self.total_size - len(indices) |
| 108 | if padding_size <= len(indices): |
| 109 | indices += indices[:padding_size] |
| 110 | else: |
| 111 | indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size] |
| 112 | else: |
| 113 | # remove tail of data to make it evenly divisible. |
| 114 | indices = indices[:self.total_size] |
| 115 | assert len(indices) == self.total_size |
| 116 | |
| 117 | # subsample |
| 118 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 119 | assert len(indices) == self.num_samples |
| 120 | |
| 121 | return iter(indices) |
| 122 | |
| 123 | def __len__(self) -> int: |
| 124 | return self.num_samples |