| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import itertools |
| from typing import Optional |
|
|
| import torch |
| import torch.distributed as dist |
| from torch.utils.data.sampler import BatchSampler as torchBatchSampler |
| from torch.utils.data.sampler import Sampler |
|
|
|
|
| class YoloBatchSampler(torchBatchSampler): |
| """ |
| This batch sampler will generate mini-batches of (mosaic, index) tuples from another sampler. |
| It works just like the :class:`torch.utils.data.sampler.BatchSampler`, |
| but it will turn on/off the mosaic aug. |
| """ |
|
|
| def __init__(self, *args, mosaic=True, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.mosaic = mosaic |
|
|
| def __iter__(self): |
| for batch in super().__iter__(): |
| yield [(self.mosaic, idx) for idx in batch] |
|
|
|
|
| class InfiniteSampler(Sampler): |
| """ |
| In training, we only care about the "infinite stream" of training data. |
| So this sampler produces an infinite stream of indices and |
| all workers cooperate to correctly shuffle the indices and sample different indices. |
| The samplers in each worker effectively produces `indices[worker_id::num_workers]` |
| where `indices` is an infinite stream of indices consisting of |
| `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True) |
| or `range(size) + range(size) + ...` (if shuffle is False) |
| """ |
|
|
| def __init__( |
| self, |
| size: int, |
| shuffle: bool = True, |
| seed: Optional[int] = 0, |
| rank=0, |
| world_size=1, |
| ): |
| """ |
| Args: |
| size (int): the total number of data of the underlying dataset to sample from |
| shuffle (bool): whether to shuffle the indices or not |
| seed (int): the initial seed of the shuffle. Must be the same |
| across all workers. If None, will use a random seed shared |
| among workers (require synchronization among all workers). |
| """ |
| self._size = size |
| assert size > 0 |
| self._shuffle = shuffle |
| self._seed = int(seed) |
|
|
| if dist.is_available() and dist.is_initialized(): |
| self._rank = dist.get_rank() |
| self._world_size = dist.get_world_size() |
| else: |
| self._rank = rank |
| self._world_size = world_size |
|
|
| def __iter__(self): |
| start = self._rank |
| yield from itertools.islice( |
| self._infinite_indices(), start, None, self._world_size |
| ) |
|
|
| def _infinite_indices(self): |
| g = torch.Generator() |
| g.manual_seed(self._seed) |
| while True: |
| if self._shuffle: |
| yield from torch.randperm(self._size, generator=g) |
| else: |
| yield from torch.arange(self._size) |
|
|
| def __len__(self): |
| return self._size // self._world_size |
|
|
