| """Unique-group batches for supervised contrastive learning.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| import random |
| from collections import defaultdict |
| from collections.abc import Iterator, Sequence |
|
|
| from torch.utils.data import Sampler |
|
|
|
|
| class UniqueGroupBatchSampler(Sampler[list[int]]): |
| def __init__( |
| self, |
| group_ids: Sequence[str], |
| groups_per_batch: int, |
| examples_per_group: int = 2, |
| seed: int = 0, |
| ) -> None: |
| if groups_per_batch < 2 or examples_per_group < 2: |
| raise ValueError("A contrastive batch needs at least two groups and two examples per group") |
| grouped: dict[str, list[int]] = defaultdict(list) |
| for index, group_id in enumerate(group_ids): |
| grouped[str(group_id)].append(index) |
| self.grouped = {key: value for key, value in grouped.items() if len(value) >= examples_per_group} |
| if len(self.grouped) < 2: |
| raise ValueError("At least two groups have enough examples") |
| self.groups_per_batch = groups_per_batch |
| self.examples_per_group = examples_per_group |
| self.seed = seed |
| self.epoch = 0 |
|
|
| def set_epoch(self, epoch: int) -> None: |
| self.epoch = int(epoch) |
|
|
| def __len__(self) -> int: |
| return math.ceil(len(self.grouped) / self.groups_per_batch) |
|
|
| def __iter__(self) -> Iterator[list[int]]: |
| random_state = random.Random(self.seed + self.epoch) |
| groups = sorted(self.grouped) |
| random_state.shuffle(groups) |
| for start in range(0, len(groups), self.groups_per_batch): |
| selected = groups[start : start + self.groups_per_batch] |
| if len(selected) < 2: |
| selected.extend(groups[: 2 - len(selected)]) |
| batch: list[int] = [] |
| for group_id in selected: |
| batch.extend(random_state.sample(self.grouped[group_id], self.examples_per_group)) |
| yield batch |
|
|
