File size: 722 Bytes
7b5e67a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
from torch.utils.data import WeightedRandomSampler
import torch
import numpy as np
class CustomWeightedRandomSampler(WeightedRandomSampler):
"""WeightedRandomSampler except allows for more than 2^24 samples to be sampled"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __iter__(self):
rand_tensor = np.random.choice(range(0, len(self.weights)),
size=self.num_samples,
p=self.weights.numpy() / torch.sum(self.weights).numpy(),
replace=self.replacement)
rand_tensor = torch.from_numpy(rand_tensor)
return iter(rand_tensor.tolist())
|