(self)
| 9 | |
| 10 | |
| 11 | def __iter__(self): |
| 12 | if self.shuffle: |
| 13 | # deterministically shuffle based on epoch and seed |
| 14 | g = torch.Generator() |
| 15 | g.manual_seed(self.seed + self.epoch) |
| 16 | indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] |
| 17 | else: |
| 18 | indices = list(range(len(self.dataset))) # type: ignore[arg-type] |
| 19 | |
| 20 | if not self.drop_last: |
| 21 | # add extra samples to make it evenly divisible |
| 22 | padding_size = self.total_size - len(indices) |
| 23 | if padding_size <= len(indices): |
| 24 | indices += indices[:padding_size] |
| 25 | else: |
| 26 | indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size] |
| 27 | else: |
| 28 | # remove tail of data to make it evenly divisible. |
| 29 | indices = indices[:self.total_size] |
| 30 | assert len(indices) == self.total_size |
| 31 | |
| 32 | # subsample |
| 33 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 34 | assert len(indices) == self.num_samples |
| 35 | |
| 36 | return iter(indices) |
nothing calls this directly
no outgoing calls
no test coverage detected