| 6 | |
| 7 | |
| 8 | class DistributedSampler(Sampler): |
| 9 | def __init__(self, dataset, world_size=None, rank=None, round_up=True): |
| 10 | if world_size is None: |
| 11 | world_size = link.get_world_size() |
| 12 | if rank is None: |
| 13 | rank = link.get_rank() |
| 14 | self.dataset = dataset |
| 15 | self.world_size = world_size |
| 16 | self.rank = rank |
| 17 | self.round_up = round_up |
| 18 | self.epoch = 0 |
| 19 | |
| 20 | self.num_samples = int( |
| 21 | math.ceil(len(self.dataset) * 1.0 / self.world_size)) |
| 22 | if self.round_up: |
| 23 | self.total_size = self.num_samples * self.world_size |
| 24 | self.length = self.num_samples |
| 25 | else: |
| 26 | self.total_size = len(self.dataset) |
| 27 | |
| 28 | if self.rank < self.world_size-1: |
| 29 | self.length = self.num_samples |
| 30 | else: |
| 31 | self.length = self.total_size - \ |
| 32 | (self.world_size-1)*self.num_samples |
| 33 | |
| 34 | def __iter__(self): |
| 35 | g = torch.Generator() |
| 36 | g.manual_seed(self.epoch) |
| 37 | indices = list(torch.randperm(len(self.dataset), generator=g)) |
| 38 | |
| 39 | if self.round_up: |
| 40 | indices += indices[:(self.total_size - len(indices))] |
| 41 | assert len(indices) == self.total_size |
| 42 | |
| 43 | offset = self.num_samples * self.rank |
| 44 | indices = indices[offset:offset + self.num_samples] |
| 45 | if self.round_up or (not self.round_up and self.rank < self.world_size-1): |
| 46 | assert len(indices) == self.num_samples |
| 47 | |
| 48 | return iter(indices) |
| 49 | |
| 50 | def __len__(self): |
| 51 | return self.length |
| 52 | |
| 53 | def set_epoch(self, epoch): |
| 54 | self.epoch = epoch |
| 55 | |
| 56 | |
| 57 | class DistributedGivenIterationSampler(Sampler): |
nothing calls this directly
no outgoing calls
no test coverage detected