(self)
| 35 | self.total_size = self.num_samples * self.num_replicas |
| 36 | |
| 37 | def __iter__(self): |
| 38 | # deterministically shuffle based on epoch |
| 39 | g = torch.Generator() |
| 40 | g.manual_seed(self.epoch) |
| 41 | indices = list(torch.randperm(len(self.dataset), generator=g)) |
| 42 | |
| 43 | # add extra samples to make it evenly divisible |
| 44 | indices += indices[:(self.total_size - len(indices))] |
| 45 | assert len(indices) == self.total_size |
| 46 | |
| 47 | # subsample |
| 48 | offset = self.num_samples * self.rank |
| 49 | indices = indices[offset:offset + self.num_samples] |
| 50 | assert len(indices) == self.num_samples |
| 51 | |
| 52 | return iter(indices) |
| 53 | |
| 54 | def __len__(self): |
| 55 | return self.num_samples |
nothing calls this directly
no outgoing calls
no test coverage detected