PyTorch [Basics] — Sampling Samplers | by Akshaj Verma | Towards Data Science

This notebook takes you through an implementation of random_split, SubsetRandomSampler, and WeightedRandomSampler on Natural Images data using PyTorch.
import numpy as npimport pandas as pdimport seaborn as snsfrom tqdm.notebook import tqdmimport matplotlib.pyplot as pltimport torchimport torchvisionimport torch.nn as nnimport torch.optim as optimimport torch.nn.functional as Ffrom torchvision import transforms, utils, datasetsfrom torch.utils.data import Dataset, DataLoader, random_split, SubsetRandomSampler, WeightedRandomSampler
Set the random seed.
np.random.seed(0)torch.manual_seed(0)
Set Seaborn style.
%matplotlib inlinesns.set_style('darkgrid')
Set the root directory for the dataset.
root_dir = "../../data/computer_vision/image_classification/natural-images/"print("The data lies here =>", root_dir)###################### OUTPUT ######################We're using => cpuThe data lies here => ../../data/computer_vision/image_classification/natural-images/
Crop the images to be of size (224, 224) and convert them to tensors.
image_transforms = {    "train": transforms.Compose([        transforms.Resize((224, 224)),        transforms.ToTensor()    ])}
Using ImageFolder, we will create our dataset. We'll only use the train folder for this blogpost.
natural_img_dataset = datasets.ImageFolder(                              root = root_dir,                              transform = image_transforms["train"]                       )natural_img_dataset
The .class_to_idx method returns the class-mapping label in the dataset.
natural_img_dataset.class_to_idx###################### OUTPUT ######################{'airplane': 0, 'car': 1, 'cat': 2, 'dog': 3, 'flower': 4, 'fruit': 5, 'motorbike': 6, 'person': 7}
We will create a dictionary called idx2class which is the reverse of class_to_idx method in PyTorch.
idx2class = {v: k for k, v in natural_img_dataset.class_to_idx.items()}idx2class###################### OUTPUT ######################{0: 'airplane', 1: 'car', 2: 'cat', 3: 'dog', 4: 'flower', 5: 'fruit', 6: 'motorbike', 7: 'person'}
To observe the distribution of different classes in a dataset object, we create a function called get_class_distribution(). This function takes a dataset as an input argument and returns a dictionary which contains the count of all classes in the dataset object.
To do this, we first initialize our count_dict where all the class counts are 0.Then we iterate over our dataset object to extract the class labels. The dataset object contains elements in the form of a tuple (x,y). So, we need to extract the item at position 1 from the tuple.Then we use the idx2class to get the class name from the class id.Finally, we update the count in our count_dict by 1 for the relevant class-key.
To do this, we first initialize our count_dict where all the class counts are 0.
Then we iterate over our dataset object to extract the class labels. The dataset object contains elements in the form of a tuple (x,y). So, we need to extract the item at position 1 from the tuple.
Then we use the idx2class to get the class name from the class id.
Finally, we update the count in our count_dict by 1 for the relevant class-key.
def get_class_distribution(dataset_obj):    count_dict = {k:0 for k,v in dataset_obj.class_to_idx.items()}        for element in dataset_obj:        y_lbl = element[1]        y_lbl = idx2class[y_lbl]        count_dict[y_lbl] += 1                return count_dictprint("Distribution of classes: \n", get_class_distribution(natural_img_dataset))###################### OUTPUT ######################Distribution of classes:  {'airplane': 727, 'car': 968, 'cat': 885, 'dog': 702, 'flower': 843, 'fruit': 1000, 'motorbike': 788, 'person': 986}
To plot our dictionary, we use the Seaborn library. We first convert our dictionary to a dataframe and then melt it. Finally, we use the function sns.barplot() to construct our plot.
plt.figure(figsize=(15,8))sns.barplot(data = pd.DataFrame.from_dict([get_class_distribution(natural_img_dataset)]).melt(), x = "variable", y="value", hue="variable").set_title('Natural Images Class Distribution')
From the above graph, we observe that the classes are imbalanced.
random_split(dataset, lengths) works directly on the dataset. The function expects 2 input arguments. The first argument is the dataset. The second is a tuple of lengths. If we want to split our dataset into 2 parts, we will provide a tuple with 2 numbers. These numbers are the sizes of the corresponding datasets after the split.
Our dataset has 6899 images. If we want to split this into 2 parts (train/test, train/val) of size (6000, 899), we will call random split as random_split(6000, 899).
Let’s split our dataset into train and val sets.
train_dataset, val_dataset = random_split(natural_img_dataset, (6000, 899))
Pass data to the dataloader.
train_loader = DataLoader(dataset=train_dataset, shuffle=True, batch_size=1)val_loader = DataLoader(dataset=val_dataset, shuffle=False, batch_size=1)print("Length of the train_loader:", len(train_loader))print("Length of the val_loader:", len(val_loader))###################### OUTPUT ######################Length of the train_loader: 6000Length of the val_loader: 899
Note that we have used a batch_size = 1. If we increase the batch_size, the number of images would be the same but the length of train/val loaders would change.
Let’s take a look at the distribution of classes in the train and val loaders.
def get_class_distribution_loaders(dataloader_obj, dataset_obj):    count_dict = {k:0 for k,v in dataset_obj.class_to_idx.items()}        for _,j in dataloader_obj:        y_idx = j.item()        y_lbl = idx2class[y_idx]        count_dict[str(y_lbl)] += 1                return count_dict
Let’s construct the plots.
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(18,7))sns.barplot(data = pd.DataFrame.from_dict([get_class_distribution_loaders(train_loader, natural_img_dataset)]).melt(), x = "variable", y="value", hue="variable",  ax=axes[0]).set_title('Train Set')sns.barplot(data = pd.DataFrame.from_dict([get_class_distribution_loaders(val_loader, natural_img_dataset)]).melt(), x = "variable", y="value", hue="variable",  ax=axes[1]).set_title('Val Set')
SubsetRandomSampler(indices) takes as input the indices of data.
We first create our samplers and then we’ll pass it to our dataloaders.
Create a list of indices.Shuffle the indices.Split the indices based on train-val percentage.Create SubsetRandomSampler.
Create a list of indices.
Shuffle the indices.
Split the indices based on train-val percentage.
Create SubsetRandomSampler.
Create a list of indices from 0 to length of dataset.
dataset_size = len(natural_img_dataset)dataset_indices = list(range(dataset_size))
Shuffle the list of indices using np.shuffle.
np.random.shuffle(dataset_indices)
Create the split index. We choose the split index to be 20% (0.2) of the dataset size.
val_split_index = int(np.floor(0.2 * dataset_size))
Slice the lists to obtain 2 lists of indices, one for train and other for test.
0-----------val_split_index------------------------------n.
Train => val_split_index to n
Val => 0 to val_split_index
train_idx, val_idx = dataset_indices[val_split_index:], dataset_indices[:val_split_index]
Finally, create samplers.
train_sampler = SubsetRandomSampler(train_idx)val_sampler = SubsetRandomSampler(val_idx)
Now, we will pass the samplers to our dataloader. Note that shuffle=True cannot be used when you're using the SubsetRandomSampler.
train_loader = DataLoader(dataset=natural_img_dataset, shuffle=False, batch_size=1, sampler=train_sampler)val_loader = DataLoader(dataset=natural_img_dataset, shuffle=False, batch_size=1, sampler=val_sampler)
Now, we’ll plot the class distribution in our dataloaders.
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(18,7))sns.barplot(data = pd.DataFrame.from_dict([get_class_distribution_loaders(train_loader, natural_img_dataset)]).melt(), x = "variable", y="value", hue="variable",  ax=axes[0]).set_title('Train Set')sns.barplot(data = pd.DataFrame.from_dict([get_class_distribution_loaders(val_loader, natural_img_dataset)]).melt(), x = "variable", y="value", hue="variable",  ax=axes[1]).set_title('Val Set')
As we can observe, the number of samples per class in the validation set is proportional to the number in train set.
WeightedRandomSampler is used, unlike random_split and SubsetRandomSampler, to ensure that each batch sees a proportional number of all classes.
Get all the target classes.Get the class weights. Class weights are the reciprocal of the number of items per class.Obtain corresponding weight for each target sample.
Get all the target classes.
Get the class weights. Class weights are the reciprocal of the number of items per class.
Obtain corresponding weight for each target sample.
Obtain the list of target classes and shuffle.
target_list = torch.tensor(natural_img_dataset.targets)
Get the class counts and calculate the weights/class by taking its reciprocal.
class_count = [i for i in get_class_distribution(natural_img_dataset).values()]class_weights = 1./torch.tensor(class_count, dtype=torch.float) class_weights###################### OUTPUT ######################tensor([0.0014, 0.0010, 0.0011, 0.0014, 0.0012, 0.0010, 0.0013, 0.0010])
Assign the weight of each class to all the samples.
class_weights_all = class_weights[target_list]class_weights_all###################### OUTPUT ######################tensor([0.0010, 0.0012, 0.0014,  ..., 0.0010, 0.0014, 0.0010])
Pass the weight and number of samples to the WeightedRandomSampler.
weighted_sampler = WeightedRandomSampler(    weights=class_weights_all,    num_samples=len(class_weights_all),    replacement=True)
Pass the sampler to the dataloader.
train_loader = DataLoader(dataset=natural_img_dataset, shuffle=False, batch_size=8, sampler=weighted_sampler)
And this is it. You can now use your dataloader to train your neural network model!
Thank you for reading. Suggestions and constructive criticism are welcome. :)
This blog post is a part of the series — “How to train you Neural Net”. You can find the series here.
You can find me on LinkedIn and Twitter. If you liked this, check out my other blogposts.