Spaces:
Sleeping
Sleeping
File size: 1,401 Bytes
20239f9 |
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 |
import torch
from torch.utils.data import Dataset
class ClassBalancedRandomSampler(torch.utils.data.Sampler):
"""
A custom sampler that sub-samples a given dataset based on class labels. Based on the RandomSampler class
This is essentially the non-ddp version of ClassBalancedDistributedSampler
Ref: https://github.com/pytorch/pytorch/blob/abe3c55a6a01c5b625eeb4fc9aab1421a5965cd2/torch/utils/data/sampler.py#L117
"""
def __init__(self, dataset: Dataset, num_samples_per_class=100, seed: int = 0) -> None:
self.dataset = dataset
self.seed = seed
# Calculate the number of samples
self.generator = torch.Generator()
self.generator.manual_seed(self.seed)
self.num_samples_per_class = num_samples_per_class
indices = dataset.generate_class_balanced_indices(self.generator,
num_samples_per_class=num_samples_per_class)
self.num_samples = len(indices)
def __iter__(self):
# Change seed for every function call
seed = int(torch.empty((), dtype=torch.int64).random_().item())
self.generator.manual_seed(seed)
indices = self.dataset.generate_class_balanced_indices(self.generator, num_samples_per_class=self.num_samples_per_class)
return iter(indices)
def __len__(self) -> int:
return self.num_samples
|