|
|
| import numpy as np
|
| from torch.utils.data.sampler import BatchSampler, Sampler
|
|
|
|
|
| class GroupedBatchSampler(BatchSampler):
|
| """
|
| Wraps another sampler to yield a mini-batch of indices.
|
| It enforces that the batch only contain elements from the same group.
|
| It also tries to provide mini-batches which follows an ordering which is
|
| as close as possible to the ordering from the original sampler.
|
| """
|
|
|
| def __init__(self, sampler, group_ids, batch_size):
|
| """
|
| Args:
|
| sampler (Sampler): Base sampler.
|
| group_ids (list[int]): If the sampler produces indices in range [0, N),
|
| `group_ids` must be a list of `N` ints which contains the group id of each sample.
|
| The group ids must be a set of integers in the range [0, num_groups).
|
| batch_size (int): Size of mini-batch.
|
| """
|
| if not isinstance(sampler, Sampler):
|
| raise ValueError(
|
| "sampler should be an instance of "
|
| "torch.utils.data.Sampler, but got sampler={}".format(sampler)
|
| )
|
| self.sampler = sampler
|
| self.group_ids = np.asarray(group_ids)
|
| assert self.group_ids.ndim == 1
|
| self.batch_size = batch_size
|
| groups = np.unique(self.group_ids).tolist()
|
|
|
|
|
| self.buffer_per_group = {k: [] for k in groups}
|
|
|
| def __iter__(self):
|
| for idx in self.sampler:
|
| group_id = self.group_ids[idx]
|
| group_buffer = self.buffer_per_group[group_id]
|
| group_buffer.append(idx)
|
| if len(group_buffer) == self.batch_size:
|
| yield group_buffer[:]
|
| del group_buffer[:]
|
|
|
| def __len__(self):
|
| raise NotImplementedError("len() of GroupedBatchSampler is not well-defined.")
|
|
|
