File size: 1,975 Bytes
d187b57 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
from torch.utils.data import Sampler
import numpy as np
import logging
from collections import defaultdict
from pathlib import Path
import torch
logger = logging.getLogger(__name__)
class MultilabelStratifiedSampler(Sampler):
def __init__(self, labels, groups, batch_size, cached_size=None):
super().__init__(None)
self.labels = np.array(labels)
self.groups = np.array(groups)
self.batch_size = batch_size
self.num_samples = len(labels)
# Simple validation
if len(self.labels) != len(self.groups):
raise ValueError("Length mismatch between labels and groups")
# Create indices per group
self.group_indices = {}
unique_groups = np.unique(self.groups)
for group in unique_groups:
indices = np.where(self.groups == group)[0]
if len(indices) > 0:
self.group_indices[group] = indices
# Calculate group probabilities
group_sizes = np.array([len(indices) for indices in self.group_indices.values()])
self.group_probs = group_sizes / group_sizes.sum()
self.valid_groups = list(self.group_indices.keys())
# Calculate number of batches
self.num_batches = self.num_samples // self.batch_size
if self.num_batches == 0:
self.num_batches = 1
self.total_samples = self.num_batches * self.batch_size
def __iter__(self):
indices = []
for _ in range(self.num_batches):
batch = []
for _ in range(self.batch_size):
# Select group and sample from it
group = np.random.choice(self.valid_groups, p=self.group_probs)
idx = np.random.choice(self.group_indices[group])
batch.append(idx)
indices.extend(batch)
return iter(indices)
def __len__(self):
return self.total_samples |