PyTorch implementations of BatchSampler that under/over sample according to a chosen parameter
alpha, in order to create a balanced training distribution.
The factory class constructs a pytorch BatchSampler to yield balanced samples from a
training distribution.
from pytorch_balanced_sampler.sampler import SamplerFactory
# which sample indices belong to each of 4 classes
class_idxs = [
[1, 2, 3, 4, 5],
[6, 7, 8, 9, 10],
[11, 12],
[13, 14, 15, 16, 17, 18, 19, 20]
]
batch_sampler = SamplerFactory().get(
class_idxs=class_idxs,
batch_size=32,
n_batches=250,
alpha=0.5,
kind='fixed'
)
dataset = Dataset( ... )
data_loader = DataLoader(dataset, batch_sampler=batch_sampler)
for data, target in data_loader:
# nice balanced batches!
...
Based on the choice of an alpha parameter in [0, 1] the sampler will adjust the sample
distribution to be between true distribution (alpha = 0), and a uniform distribution
(alpha = 1).
Overrepresented classes will be undersampled, and underrepresented classes oversampled. Here's an example from an imbalanced data distribution I was working with a while ago:
If you select kind='fixed', each batch generated will contain a consistent proportion of
classes. Eg. if we have 5 classes, we might receive batches like:
Batch: 0
Classes: [1, 0, 0, 0, 2, 4, 0, 2, 0, 0, 3, 2, 1, 0, 2, 0, 0, 3, 0, 0, 4, 4, 0, 2, 1, 3, 3, 1, 2, 0, 0, 4]
Counts: {0: 14, 1: 4, 2: 6, 3: 4, 4: 4}
Batch: 1
Classes: [4, 1, 1, 2, 0, 0, 0, 4, 2, 4, 0, 3, 1, 3, 0, 0, 3, 2, 0, 2, 4, 2, 0, 0, 2, 3, 0, 1, 0, 0, 0, 0]
Counts: {0: 14, 1: 4, 2: 6, 3: 4, 4: 4}
Batch: 2
Classes: [0, 4, 0, 0, 0, 3, 3, 2, 0, 4, 2, 3, 0, 3, 2, 0, 0, 1, 2, 2, 0, 1, 0, 0, 4, 0, 2, 1, 1, 4, 0, 0]
Counts: {0: 14, 1: 4, 2: 6, 3: 4, 4: 4}
Note that the class counts are the same for each batch.
If you don't want to fix the number of each class in each batch, you can select kind='random',
which will use sampling with replacement. The samples will be weighted as to produce the target
class distribution on average.
pytorch_balanced_sampler was written by Karl Hornlund.