(self)
| 28 | self.total_size = self.num_samples * self.num_replicas |
| 29 | |
| 30 | def __iter__(self): |
| 31 | # deterministically shuffle based on epoch |
| 32 | g = torch.Generator() |
| 33 | g.manual_seed(self.epoch) |
| 34 | indices = torch.randperm(self.total_size, generator=g).tolist() |
| 35 | |
| 36 | dataset_size = len(self.dataset) |
| 37 | indices = [v % dataset_size for v in indices] |
| 38 | |
| 39 | # subsample |
| 40 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 41 | assert len(indices) == self.num_samples |
| 42 | |
| 43 | return iter(indices) |
| 44 | |
| 45 | def __len__(self): |
| 46 | return self.num_samples |
nothing calls this directly
no outgoing calls
no test coverage detected