Using Weighted Random Sampler in PyTorch

Sometimes there are scenarios where you have way lesser number of samples for some of the classes where as other classes have lots of samples. In such a scenario, you don’t want a training batch to be contain samples just from a few of the classes with lots of samples. Ideally, a training batch should contain represent a good spread of the dataset. In PyTorch this can be achieved using a weighted random sampler.

In this short post, I will walk you through the process of creating a random weighted sampler in PyTorch.

To start off, lets assume you have a dataset with images grouped in folders based on their class. We can use a ImageFolder to create a dataset from it.

from torchvision import datasets

dataset = datasets.ImageFolder(root=data_dir, transform=image_transforms)

Also, lets split this dataset into training and validation sets,

dataset_size = dataset.__len__()
train_count = int(dataset_size * 0.7)
val_count = dataset_size - train_count
train_dataset, valid_dataset = random_split(dataset, [train_count, val_count])

Note: I am taking a ImageFolder and training/validation splits just to emulate a real world example. You can work with any pytorch dataset.

We will be using a weighted random sampler just for the training set. For validation set, we don’t care about balancing a batch. Now that we have the train_dataset, you need to define the weights for each class which would be inversely proportional to the number of samples for each class.

First, lets find the number of samples for each class.

import numpy as np 

y_train_indices = train_dataset.indices

y_train = [dataset.targets[i] for i in y_train_indices]

class_sample_count = np.array(
    [len(np.where(y_train == t)[0]) for t in np.unique(y_train)])

Next, we need to find the weights for each class.

weight = 1. / class_sample_count
samples_weight = np.array([weight[t] for t in y_train])
samples_weight = torch.from_numpy(samples_weight)

Now, that we have the weights for each of the classes, we can define a sampler.

sampler = WeightedRandomSampler(samples_weight.type('torch.DoubleTensor'), len(samples_weight))

Finally, we can use the sampler, while defining the Dataloader.

train_dataloader = DataLoader(train_dataset, batch_size=4, sampler=sampler)

Thats it for this post. I hope you found this post useful.