File size: 3,906 Bytes
212111c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
"""Implementation of a bucketed data sampler from PyTorch-NLP.
Modified by Roshan Rao.
See https://github.com/PetrochukM/PyTorch-NLP/
"""
import typing
import math
import operator
from torch.utils.data.sampler import Sampler
from torch.utils.data.sampler import BatchSampler
from torch.utils.data.sampler import SubsetRandomSampler
class SortedSampler(Sampler):
""" Samples elements sequentially, always in the same order.
Args:
data (iterable): Iterable data.
sort_key (callable): Specifies a function of one argument that is used to extract a
numerical comparison key from each list element.
Example:
>>> list(SortedSampler(range(10), sort_key=lambda i: -i))
[9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
"""
def __init__(self,
dataset,
sort_key: typing.Callable[[int], typing.Any],
indices: typing.Optional[typing.Iterable[int]] = None):
super().__init__(dataset)
self.dataset = dataset
self.sort_key = sort_key
if indices is None:
sort_keys = map(sort_key, dataset)
else:
sort_keys = ((i, sort_key(dataset[i])) for i in indices)
self.sorted_indices = [i for i, _ in sorted(sort_keys, key=operator.itemgetter(1))]
def __iter__(self):
return iter(self.sorted_indices)
def __len__(self):
return len(self.dataset)
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. Provides ~10-25 percent speedup.
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.
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 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,
dataset,
bucket_size_multiplier=100):
super().__init__(sampler, batch_size, drop_last)
self.sort_key = sort_key
self.dataset = dataset
self.bucket_sampler = BatchSampler(
sampler, min(batch_size * bucket_size_multiplier, len(sampler)), False)
def __iter__(self):
for bucket in self.bucket_sampler:
sorted_sampler = SortedSampler(self.dataset, self.sort_key, indices=bucket)
for batch in SubsetRandomSampler(
list(BatchSampler(sorted_sampler, self.batch_size, self.drop_last))):
yield 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)
|