(self)
| 30 | self.total_size = len(self.dataset) |
| 31 | |
| 32 | def __iter__(self): |
| 33 | # deterministically shuffle based on epoch |
| 34 | if self.shuffle: |
| 35 | g = torch.Generator() |
| 36 | g.manual_seed(self.epoch) |
| 37 | indices = torch.randperm(len(self.dataset), generator=g).tolist() |
| 38 | else: |
| 39 | indices = torch.arange(len(self.dataset)).tolist() |
| 40 | |
| 41 | # add extra samples to make it evenly divisible |
| 42 | if self.round_up: |
| 43 | indices = ( |
| 44 | indices * |
| 45 | int(self.total_size / len(indices) + 1))[:self.total_size] |
| 46 | assert len(indices) == self.total_size |
| 47 | |
| 48 | # subsample |
| 49 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 50 | if self.round_up: |
| 51 | assert len(indices) == self.num_samples |
| 52 | |
| 53 | return iter(indices) |
| 54 | |
| 55 | |
| 56 | def build_dataloader(dataset: Dataset, |
nothing calls this directly
no outgoing calls
no test coverage detected