| |
| """Hard-negative batch sampler based on confusion pairs. |
| |
| Strategy: each batch contains 50% normal class-balanced samples + |
| 50% samples whose class is in a confusion pair with another. This |
| forces the loss to separate hardest pairs. |
| |
| Usage: |
| from hard_neg_sampler import ConfusionBatchSampler |
| sampler = ConfusionBatchSampler(labels, confusion_pairs, batch_size=64) |
| loader = DataLoader(ds, batch_sampler=sampler, ...) |
| """ |
| import random |
| from collections import defaultdict |
| import numpy as np |
| import torch |
| from torch.utils.data import Sampler |
|
|
| class ConfusionBatchSampler(Sampler): |
| """Yields batches weighted towards confusion pairs.""" |
| def __init__(self, labels, confusion_pairs, batch_size=64, hard_frac=0.5, |
| seed=0, num_batches=None): |
| self.labels = np.asarray(labels) |
| self.batch_size = batch_size |
| self.hard_frac = hard_frac |
| |
| self.idx_by_class = defaultdict(list) |
| for i, y in enumerate(self.labels): |
| self.idx_by_class[int(y)].append(i) |
| |
| self.hard_classes = set() |
| for a, b in confusion_pairs: |
| self.hard_classes.add(int(a)); self.hard_classes.add(int(b)) |
| self.hard_classes = [c for c in self.hard_classes if self.idx_by_class[c]] |
| self.all_classes = list(self.idx_by_class) |
| self.rng = random.Random(seed) |
| self.num_batches = num_batches or (len(self.labels) // batch_size) |
|
|
| def __iter__(self): |
| for _ in range(self.num_batches): |
| n_hard = int(self.batch_size * self.hard_frac) |
| batch = [] |
| for _ in range(n_hard): |
| if not self.hard_classes: break |
| c = self.rng.choice(self.hard_classes) |
| batch.append(self.rng.choice(self.idx_by_class[c])) |
| while len(batch) < self.batch_size: |
| c = self.rng.choice(self.all_classes) |
| batch.append(self.rng.choice(self.idx_by_class[c])) |
| yield batch |
|
|
| def __len__(self): return self.num_batches |
|
|
| def extract_confusion_pairs(probs, targets, top_n=50): |
| """From val probs+targets, find top-N confused (true, pred) pairs.""" |
| from collections import Counter |
| pred = probs.argmax(-1) |
| c = Counter() |
| for i in range(len(targets)): |
| t, p = int(targets[i]), int(pred[i]) |
| if t != p: c[tuple(sorted((t, p)))] += 1 |
| return [list(k) for k, _ in c.most_common(top_n)] |
|
|
| if __name__ == '__main__': |
| labels = np.random.randint(0, 10, 200) |
| pairs = [(0, 1), (2, 3), (4, 5)] |
| s = ConfusionBatchSampler(labels, pairs, batch_size=16, num_batches=4) |
| for b in s: print(f"batch: {b}") |
|
|
