Wraps another sampler to yield a mini-batch of indices. Args: sampler (Sampler): Base sampler. batch_size (int): Size of mini-batch. drop_last (bool): If ``True``, the sampler will drop the last batch if its size would be less than ``batch_size`` Example
| 94 | |
| 95 | |
| 96 | class BatchSampler(object): |
| 97 | """Wraps another sampler to yield a mini-batch of indices. |
| 98 | |
| 99 | Args: |
| 100 | sampler (Sampler): Base sampler. |
| 101 | batch_size (int): Size of mini-batch. |
| 102 | drop_last (bool): If ``True``, the sampler will drop the last batch if |
| 103 | its size would be less than ``batch_size`` |
| 104 | |
| 105 | Example: |
| 106 | >>> list(BatchSampler(range(10), batch_size=3, drop_last=False)) |
| 107 | [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] |
| 108 | >>> list(BatchSampler(range(10), batch_size=3, drop_last=True)) |
| 109 | [[0, 1, 2], [3, 4, 5], [6, 7, 8]] |
| 110 | """ |
| 111 | |
| 112 | def __init__(self, sampler, batch_size, drop_last): |
| 113 | self.sampler = sampler |
| 114 | self.batch_size = batch_size |
| 115 | self.drop_last = drop_last |
| 116 | |
| 117 | def __iter__(self): |
| 118 | batch = [] |
| 119 | for idx in self.sampler: |
| 120 | batch.append(idx) |
| 121 | if len(batch) == self.batch_size: |
| 122 | yield batch |
| 123 | batch = [] |
| 124 | if len(batch) > 0 and not self.drop_last: |
| 125 | yield batch |
| 126 | |
| 127 | def __len__(self): |
| 128 | if self.drop_last: |
| 129 | return len(self.sampler) // self.batch_size |
| 130 | else: |
| 131 | return (len(self.sampler) + self.batch_size - 1) // self.batch_size |