| |
| 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.") |
|
|
