Produce indices for inference. Inference needs to run on the __exact__ set of samples, therefore when the total number of samples is not divisible by the number of workers, this sampler produces different number of samples on different workers.
| 171 | |
| 172 | |
| 173 | class InferenceSampler(Sampler): |
| 174 | """ |
| 175 | Produce indices for inference. |
| 176 | Inference needs to run on the __exact__ set of samples, |
| 177 | therefore when the total number of samples is not divisible by the number of workers, |
| 178 | this sampler produces different number of samples on different workers. |
| 179 | """ |
| 180 | |
| 181 | def __init__(self, size: int): |
| 182 | """ |
| 183 | Args: |
| 184 | size (int): the total number of data of the underlying dataset to sample from |
| 185 | """ |
| 186 | self._size = size |
| 187 | assert size > 0 |
| 188 | self._rank = comm.get_rank() |
| 189 | self._world_size = comm.get_world_size() |
| 190 | |
| 191 | shard_size = (self._size - 1) // self._world_size + 1 |
| 192 | begin = shard_size * self._rank |
| 193 | end = min(shard_size * (self._rank + 1), self._size) |
| 194 | self._local_indices = range(begin, end) |
| 195 | |
| 196 | def __iter__(self): |
| 197 | yield from self._local_indices |
| 198 | |
| 199 | def __len__(self): |
| 200 | return len(self._local_indices) |
no outgoing calls
no test coverage detected