import math
from torch.utils.data.sampler import BatchSampler
from torch.utils.data.sampler import SubsetRandomSampler
from torchnlp.samplers.sorted_sampler import SortedSampler
from torchnlp.utils import identity
[docs]class BucketBatchSampler(BatchSampler):
""" `BucketBatchSampler` toggles between `sampler` batches and sorted batches.
Typically, the `sampler` will be a `RandomSampler` allowing the user to toggle between
random batches and sorted batches. A larger `bucket_size_multiplier` is more sorted and vice
versa.
Background:
``BucketBatchSampler`` is similar to a ``BucketIterator`` found in popular libraries like
``AllenNLP`` and ``torchtext``. A ``BucketIterator`` pools together examples with a similar
size length to reduce the padding required for each batch while maintaining some noise
through bucketing.
**AllenNLP Implementation:**
https://github.com/allenai/allennlp/blob/master/allennlp/data/iterators/bucket_iterator.py
**torchtext Implementation:**
https://github.com/pytorch/text/blob/master/torchtext/data/iterator.py#L225
Args:
sampler (torch.data.utils.sampler.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`.
sort_key (callable, optional): Callable to specify a comparison key for sorting.
bucket_size_multiplier (int, optional): Buckets are of size
`batch_size * bucket_size_multiplier`.
Example:
>>> from torchnlp.random import set_seed
>>> set_seed(123)
>>>
>>> from torch.utils.data.sampler import SequentialSampler
>>> sampler = SequentialSampler(list(range(10)))
>>> list(BucketBatchSampler(sampler, batch_size=3, drop_last=False))
[[6, 7, 8], [0, 1, 2], [3, 4, 5], [9]]
>>> list(BucketBatchSampler(sampler, batch_size=3, drop_last=True))
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]
"""
def __init__(self,
sampler,
batch_size,
drop_last,
sort_key=identity,
bucket_size_multiplier=100):
super().__init__(sampler, batch_size, drop_last)
self.sort_key = sort_key
_bucket_size = batch_size * bucket_size_multiplier
if hasattr(sampler, "__len__"):
_bucket_size = min(_bucket_size, len(sampler))
self.bucket_sampler = BatchSampler(sampler, _bucket_size, False)
def __iter__(self):
for bucket in self.bucket_sampler:
sorted_sampler = SortedSampler(bucket, self.sort_key)
for batch in SubsetRandomSampler(
list(BatchSampler(sorted_sampler, self.batch_size, self.drop_last))):
yield [bucket[i] for i in batch]
def __len__(self):
if self.drop_last:
return len(self.sampler) // self.batch_size
else:
return math.ceil(len(self.sampler) / self.batch_size)