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