| 116 | data samplers (sequential, random, WeightedRandomSampler, etc.) with this batch sampler. |
| 117 | """ |
| 118 | def __init__(self, sampler, batch_size, drop_last, rank=-1, world_size=2, wrap_last=False, gradient_accumulation_steps=None): |
| 119 | super(DistributedBatchSampler, self).__init__(sampler, batch_size, drop_last) |
| 120 | if rank == -1: |
| 121 | assert False, 'should not be here' |
| 122 | self.rank = rank |
| 123 | self.world_size = world_size |
| 124 | self.sampler.wrap_around = 0 |
| 125 | self.wrap_around = 0 |
| 126 | self.wrap_last = wrap_last |
| 127 | self.start_iter = 0 |
| 128 | self.effective_batch_size = batch_size if gradient_accumulation_steps is None else batch_size * gradient_accumulation_steps |
| 129 | |
| 130 | def __iter__(self): |
| 131 | batch = [] |