BucketsSampler
| 105 | |
| 106 | |
| 107 | class BucketsSampler(object): |
| 108 | """BucketsSampler""" |
| 109 | def __init__(self, buckets, batch_size, shuffle=False): |
| 110 | self.batch_size = batch_size |
| 111 | self.shuffle = shuffle |
| 112 | self.sizes, self.buckets = zip(*[(size, bucket) for size, bucket in buckets.items()]) |
| 113 | # the number of chunks in each bucket, which is clipped by range [1, len(bucket)] |
| 114 | self.chunks = [] |
| 115 | for size, bucket in zip(self.sizes, self.buckets): |
| 116 | max_ch = max(math.ceil(size * len(bucket) / batch_size), 1) |
| 117 | chunk = min(len(bucket), int(max_ch)) |
| 118 | self.chunks.append(chunk) |
| 119 | |
| 120 | def __iter__(self): |
| 121 | """Returns an iterator, randomly or sequentially returns a batch id""" |
| 122 | range_fn = np.random.permutation if self.shuffle else np.arange |
| 123 | for i in range_fn(len(self.buckets)).tolist(): |
| 124 | split_sizes = [(len(self.buckets[i]) - j - 1) // self.chunks[i] + 1 for j in range(self.chunks[i])] |
| 125 | for batch in np.split(range_fn(len(self.buckets[i])), np.cumsum(split_sizes)): |
| 126 | if len(batch): |
| 127 | yield [self.buckets[i][j] for j in batch.tolist()] |
| 128 | |
| 129 | def __len__(self): |
| 130 | """Returns the number of batches""" |
| 131 | return sum(self.chunks) |
| 132 | |
| 133 | |
| 134 | class SequentialSampler(object): |