This batch sampler will generate mini-batches of (dim, index) tuples from another sampler. It works just like the :class:`torch.utils.data.sampler.BatchSampler`, but it will prepend a dimension, whilst ensuring it stays the same across one mini-batch.
| 12 | |
| 13 | |
| 14 | class YoloBatchSampler(torchBatchSampler): |
| 15 | """ |
| 16 | This batch sampler will generate mini-batches of (dim, index) tuples from another sampler. |
| 17 | It works just like the :class:`torch.utils.data.sampler.BatchSampler`, |
| 18 | but it will prepend a dimension, whilst ensuring it stays the same across one mini-batch. |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, *args, input_dimension=None, mosaic=True, **kwargs): |
| 22 | super().__init__(*args, **kwargs) |
| 23 | self.input_dim = input_dimension |
| 24 | self.new_input_dim = None |
| 25 | self.mosaic = mosaic |
| 26 | |
| 27 | def __iter__(self): |
| 28 | self.__set_input_dim() |
| 29 | for batch in super().__iter__(): |
| 30 | yield [(self.input_dim, idx, self.mosaic) for idx in batch] |
| 31 | self.__set_input_dim() |
| 32 | |
| 33 | def __set_input_dim(self): |
| 34 | """ This function randomly changes the the input dimension of the dataset. """ |
| 35 | if self.new_input_dim is not None: |
| 36 | self.input_dim = (self.new_input_dim[0], self.new_input_dim[1]) |
| 37 | self.new_input_dim = None |
| 38 | |
| 39 | |
| 40 | class InfiniteSampler(Sampler): |
no outgoing calls
no test coverage detected