| 3 | |
| 4 | |
| 5 | class DistributedSampler(_DistributedSampler): |
| 6 | def __init__(self, |
| 7 | dataset, |
| 8 | num_replicas=None, |
| 9 | rank=None, |
| 10 | shuffle=True, |
| 11 | round_up=True): |
| 12 | super().__init__(dataset, num_replicas=num_replicas, rank=rank) |
| 13 | self.shuffle = shuffle |
| 14 | self.round_up = round_up |
| 15 | if self.round_up: |
| 16 | self.total_size = self.num_samples * self.num_replicas |
| 17 | else: |
| 18 | self.total_size = len(self.dataset) |
| 19 | |
| 20 | def __iter__(self): |
| 21 | # deterministically shuffle based on epoch |
| 22 | if self.shuffle: |
| 23 | g = torch.Generator() |
| 24 | g.manual_seed(self.epoch) |
| 25 | indices = torch.randperm(len(self.dataset), generator=g).tolist() |
| 26 | else: |
| 27 | indices = torch.arange(len(self.dataset)).tolist() |
| 28 | |
| 29 | # add extra samples to make it evenly divisible |
| 30 | if self.round_up: |
| 31 | indices = ( |
| 32 | indices * |
| 33 | int(self.total_size / len(indices) + 1))[:self.total_size] |
| 34 | assert len(indices) == self.total_size |
| 35 | |
| 36 | # subsample |
| 37 | indices = indices[self.rank:self.total_size:self.num_replicas] |
| 38 | if self.round_up: |
| 39 | assert len(indices) == self.num_samples |
| 40 | |
| 41 | return iter(indices) |