Similar to normal implementation of distributed sampler, except implementation is at the batch sampler level, instead of just the sampler level. This allows wrapping of arbitrary data samplers (sequential, random, WeightedRandomSampler, etc.) with this batch sampler. The `in
| 76 | |
| 77 | |
| 78 | class DistributedBatchSampler(data.sampler.BatchSampler): |
| 79 | """Similar to normal implementation of distributed sampler, except |
| 80 | implementation is at the batch sampler level, instead of just the |
| 81 | sampler level. This allows wrapping of arbitrary data samplers |
| 82 | (sequential, random, WeightedRandomSampler, etc.) with this batch |
| 83 | sampler. |
| 84 | |
| 85 | The `interleave` argument specifies how to distribute a batch. A value |
| 86 | of True combined with the above random sampler is equivalent to pytorch's |
| 87 | torch.utils.data.distributed.DistributedSampler. |
| 88 | |
| 89 | For the following batch [0,1,2,3,4,5,6,7] and data parallelism of 2 |
| 90 | specifying True will result in the following samples for each gpu: |
| 91 | GPU0: [0,2,4,6] GPU1: [1,3,5,7] |
| 92 | specifying False will result in the following samples: |
| 93 | GPU0: [0,1,2,3] GPU1: [4,5,6,7]""" |
| 94 | |
| 95 | def __init__(self, sampler, batch_size, drop_last, rank=-1, |
| 96 | world_size=2, wrap_last=False, interleave=False): |
| 97 | super(DistributedBatchSampler, self).__init__(sampler, batch_size, |
| 98 | drop_last) |
| 99 | if rank == -1: |
| 100 | assert False, 'should not be here' |
| 101 | rank = torch.distributed.get_rank() |
| 102 | self.rank = rank |
| 103 | self.world_size = world_size |
| 104 | self.sampler.wrap_around = 0 |
| 105 | self.wrap_around = 0 |
| 106 | self.wrap_last = wrap_last |
| 107 | self.start_iter = 0 |
| 108 | self.interleave = interleave |
| 109 | |
| 110 | def __iter__(self): |
| 111 | batch = [] |
| 112 | i = 0 |
| 113 | for idx in self.data_iterator(self.sampler, wrap_around=False): |
| 114 | batch.append(idx) |
| 115 | if len(batch) == self.batch_size: |
| 116 | tbatch = self._batch(batch) |
| 117 | if i >= self.start_iter: |
| 118 | yield tbatch |
| 119 | self.start_iter = 0 |
| 120 | i += 1 |
| 121 | batch = [] |
| 122 | batch_len = len(batch) |
| 123 | if batch_len > 0 and not self.drop_last: |
| 124 | if self.wrap_last: |
| 125 | self.sampler.wrap_around -= (self.batch_size) |
| 126 | self.wrap_around += (len(batch)) |
| 127 | self.wrap_around %= self.batch_size |
| 128 | yield self._batch(batch) |
| 129 | if self.wrap_last: |
| 130 | self.sampler.wrap_around += self.batch_size |
| 131 | |
| 132 | def data_iterator(self, _iter, wrap_around=False): |
| 133 | """iterates through data and handles wrap around""" |
| 134 | for i, idx in enumerate(_iter): |
| 135 | if i < self.wrap_around % self.batch_size: |
no outgoing calls
no test coverage detected