Forcing Ration on Batches in PyTorch DataLoader
I came across a binary classifcation task where the data is heavily imbalanced. (I'm looking at 80:1)
Through undersampling, the data ratio is at now 20:1.\
Now, the undersampled/treated data is loaded in to dataloader as below. (this is an nlp task)
train_inputs = torch.tensor(input_ids)
train_labels = torch.tensor(labels)
train_masks = torch.tensor(attention_masks)

train_data = TensorDataset(train_inputs, train_masks, train_labels)
if is_distributed:
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
else:
    train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)

For each batch, I want to make sure/enforce that each batch has the ration of 20:1 in terms of each classifications. Is there a pytorch builtin method that allows me to enforce such conditions ?

You can use WeightedRandomSampler with replacement set to true.
Just multiply the weight of the positive examples by 20 to bias the ratio towards 20:1. (assuming positives are 20 times more than negatives here)
# labels is a numpy array of shape n,1 containing 1 and 0 for each datapoint
weights = np.ones(labels.shape)
weights[labels==1] *= 20
# now the samples will be drawn with a 20:1 ratio
sampler = WeightedRandomSampler(weights=weights, num_samples=len(labels), replacement=True)

Note that setting replacement to True is necessary to enforce each batch to be sampled with the ratio but you might see an example twice during training.
However if the ratio is naturally occurring in your dataset you can leave it to False.
IMHO:
Random sampling do not ensure a 20:1 ratio for each batch but only for an expected value of the ratio (on average the ratio of the batch(sample) will tend to the ratio of the dataset(population)) so it should not impact the average gradient as well during training (in theory). But I see how in practice you might want to have more control.
