| 81 | |
| 82 | class DistributedSequentialSampler(data.sampler.Sampler): |
| 83 | def __init__(self, num_samples, train_iters, batch_size, rank=-1, world_size=2): |
| 84 | super().__init__(num_samples) |
| 85 | if rank == -1: |
| 86 | rank = 0 |
| 87 | world_size = 1 |
| 88 | self.num_samples = num_samples |
| 89 | self.rank = rank |
| 90 | self.world_size = world_size |
| 91 | self.start_iter = 0 |
| 92 | self.train_iters = train_iters |
| 93 | self.batch_size = batch_size |
| 94 | self.batch_bias = [i * (num_samples // batch_size) for i in range(batch_size)] |
| 95 | |
| 96 | def __iter__(self): |
| 97 | for idx in range(self.start_iter, self.train_iters * 10): |