| 55 | |
| 56 | |
| 57 | class DistributedGivenIterationSampler(Sampler): |
| 58 | def __init__(self, dataset, total_iter, batch_size, world_size=None, rank=None, last_iter=0): |
| 59 | if world_size is None: |
| 60 | world_size = link.get_world_size() |
| 61 | if rank is None: |
| 62 | rank = link.get_rank() |
| 63 | assert rank < world_size |
| 64 | self.dataset = dataset |
| 65 | self.total_iter = total_iter |
| 66 | self.batch_size = batch_size |
| 67 | self.world_size = world_size |
| 68 | self.rank = rank |
| 69 | self.last_iter = last_iter |
| 70 | |
| 71 | self.total_size = self.total_iter*self.batch_size |
| 72 | |
| 73 | self.indices = self.gen_new_list() |
| 74 | self.call = 0 |
| 75 | |
| 76 | def __iter__(self): |
| 77 | if self.call == 0: |
| 78 | self.call = 1 |
| 79 | return iter(self.indices[self.last_iter*self.batch_size:]) |
| 80 | else: |
| 81 | raise RuntimeError( |
| 82 | "this sampler is not designed to be called more than once!!") |
| 83 | |
| 84 | def gen_new_list(self): |
| 85 | np.random.seed(0) |
| 86 | all_size = self.total_size * self.world_size |
| 87 | indices = np.arange(len(self.dataset)) |
| 88 | indices = indices[:all_size] |
| 89 | num_repeat = (all_size-1) // indices.shape[0] + 1 |
| 90 | |
| 91 | indices = np.tile(indices, num_repeat) |
| 92 | indices = indices[:all_size] |
| 93 | |
| 94 | np.random.shuffle(indices) |
| 95 | beg = self.total_size * self.rank |
| 96 | indices = indices[beg:beg+self.total_size] |
| 97 | |
| 98 | assert len(indices) == self.total_size |
| 99 | |
| 100 | return indices |
| 101 | |
| 102 | def __len__(self): |
| 103 | # note here we do not take last iter into consideration, since __len__ |
| 104 | # should only be used for displaying, the correct remaining size is |
| 105 | # handled by dataloader |
| 106 | return self.total_size |
| 107 | |
| 108 | |
| 109 | class DistributedEpochSampler(Sampler): |
nothing calls this directly
no outgoing calls
no test coverage detected