Custom dataset in Pytorch —Part 1. Images | by Utkarsh Garg | Towards Data Science

Pytorch has a great ecosystem to load custom datasets for training machine learning models. This is the first part of the two-part series on loading Custom Datasets in Pytorch. In Part 2 we’ll explore loading a custom dataset for a Machine Translation task. In this walkthrough, we’ll learn how to load a custom image dataset for classification. The code for this walkthrough can also be found on Github.
We’ll be using the landmark dataset available here. It is a subset of the Google Landmark Data v2. It has 50 classes and contains various landmarks from around the globe.
For this exercise, we’ll keep the following folder structure:
This is a straightforward folder structure with a root folder as the Train/Test folders containing classes with images inside them. As we’ll see, it doesn’t matter in what structure we get the data in. The data can all be in a single folder with class names in the image names (like “Cat_001.jpg”) or even in a CSV, we can process all this in our custom dataset class.
Furthermore, we’ll be using Albumentations library for image augmentation. This library contains a huge number of available options for image augmentations. So, we’ll be learning about how to use it in our custom dataset pipeline. You can install it using :
pip install -U albumentations
Next, we define our augmentations using Albumentations. We define different augmentations for train and test. We apply transformations related to crop/rotation, colour/saturation, and brightness on training data. We also normalise both train and test data with image net mean and std deviation. Finally, we convert the data to PyTorch tensor using ToTensor().
Next, we create the Train, Valid, and Test sets. Here we create separate lists of image paths for Train, Valid, and Test sets. These will be used in our Dataset class which will be defined for a custom dataset.
We get the following output:
train_image_path example:  images/train/15.Central_Park/462f876f97d424a2.jpgclass example:  42.Death_Valley_National_ParkTrain size: 3996Valid size: 1000Test size: 1250
We can’t use the class names directly for models. We create mappings of classes to index and index to classes.
We get an output for idx_to_class like this:
{0: '42.Death_Valley_National_Park', 1: '39.Gateway_of_India', 2: '13.Yellowstone_National_Park', 3: '44.Trevi_Fountain', 4: '32.Hanging_Temple'}
This is the core of our custom dataset. The structure of the dataset class is something like this:
We create our LandmarkDataset class by inheriting the Dataset class:
from torch.utils.data import Dataset
First, we define the __init__ function. As soon as we create an instance of our LandMarkDataset class, this function is called by default. This function should contain all operations that we want to run on the whole dataset (eg. train) once. The usage of this will be more clear in the next part of this series where we create a custom machine translation dataset. For now, we define the variables for image_paths and transforms for the corresponding Train, Valid, and Test sets.
Then we have the __len__ function which just returns the length of the dataset. This is used afterward by the DataLoader to create batches.
And finally, we have __getitem__. This processes and returns 1 datapoint at a time.
As can be seen above, __getitem__ expects an index. This is handled automatically by the dataloader which for every image in the batch runs __getitem__. In the code for __getitem__, we load the image at index “idx”, extract the label from the file path and then run it through our defined transform. The function returns the Tensor of the image array and its corresponding label.
After creating the train_dataset, we can access one example as follows:
output:The shape of tensor for 50th image in train dataset:  torch.Size([3, 256, 256])The label for 50th image in train dataset: 37
Let’s visualize some images after augmentation through the train_dataset.
Some more,
The final step. DataLoader class is used to load data in batches for the model. This helps us processing data in mini-batches that can fit within our GPU’s RAM. First, we import the DataLoader:
from torch.utils.data import DataLoader
Initiating the dataloader by sending in an object of the dataset and the batch size.
Once we have the dataloader instance — train_loader, we can use an iterator to access the data like this:
#batch of image tensornext(iter(train_loader))[0].shapeoutput:torch.Size([64, 3, 256, 256])#batch of the corresponding labelsnext(iter(train_loader))[1].shapeoutput: torch.Size([64])
This is what we use to batch out the data in our training loop. Every time we run the iterator, the dataloader selects the next 64 indexes and runs it through the __getitem__ in dataset class one by one and then returns it to the training loop.
In part 1 of this 2 part series, we saw how we can write our own custom data pipeline. We also learned to use Albumentations for image augmentation. Further, we understood how Dataset and Dataloader classes work internally. In the next part, we’ll up the level by creating a custom dataset class for a Machine Translation task.
See you in the next one!
Part 2 can be found here.