Dataset & DataLoader

Dataset & DataLoader#

Code for processing data can get messy and hard to maintain. We ideally want our dataset code to be decoupled from our model training code for better readability and modularity. PyTorch provides two data primitives to help us manage data in a way that is easy to work with.

  • Dataset represents a collection of data samples and their corresponding labels.

  • DataLoader wraps an iterable around a Dataset object to enable batching, shuffling, etc.

PyTorch domain libraries provide a number of pre-loaded datasets that subclass Dataset. The TorchVision library specifically includes datasets for many computer vision tasks, such as MNIST, ImageNet, COCO, and much more. In this tutorial, we will use the MNIST dataset.

Hide code cell source

 
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import v2
import matplotlib.pyplot as plt

Loading the dataset#

MNIST provides official training and test partitions. We split the training partition again and keep the test partition for final evaluation. This gives us three sets:

  • the training set fits the parameters,

  • the validation set guides choices such as architecture and training duration,

  • the test set is reserved for one final evaluation.

We load the data with the following parameters:

  • root is the path where the train/test data is stored,

  • train specifies training or test partition,

  • download=True downloads the data from the internet if it’s not available at root,

  • transform accepts a function that transforms the images in the dataset as a preprocessing step.

preprocess = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])

full_ds = MNIST('.data', train=True, download=True, transform=preprocess)
test_ds  = MNIST('.data', train=False, download=True, transform=preprocess)

train_ds, valid_ds = torch.utils.data.random_split(
    full_ds, [0.9, 0.1], generator=torch.Generator().manual_seed(42)
)

print('Train set:', len(train_ds))
print('Valid set:', len(valid_ds))
print('Test  set:', len(test_ds))

Train set: 54000
Valid set: 6000
Test  set: 10000

The variables train_ds, valid_ds, and test_ds are Dataset objects that contain samples represented as tuples of images and labels. The train_ds is used for training the model, valid_ds is used for validating the model during training, and test_ds is used for evaluating the model after training. We can inspect valid_ds repeatedly while developing the model, but we leave test_ds untouched until all choices are final.

Note

  • MNIST is roughly balanced, so a random split is adequate.

  • For an imbalanced dataset, use a stratified split.

  • When samples share a person, object, or source, split by that group to avoid leakage between the training set and the validation/test sets.

Preprocessing#

Data rarely arrives in the exact representation a model expects. Loading a sample from the MNIST Dataset returns a PIL image and an integer from 0 to 9. Here we convert images to floating-point tensors and rescale integer pixels from [0, 255] to [0,1]. This is not statistical standardization, which would additionally subtract a mean and divide by a standard deviation.

Hereabove, the MNIST Dataset was instructed to preprocess the images with the following transforms.

  • ToImage() - Convert a PIL Image, a NumPy array, or a PyTorch tensor to the Image type, which is a subclass of torch.Tensor defined in TorchVision to facilitate image processing.

  • ToDtype() - Convert the input values to floats, and optionally rescale them to the range [0, 1].

Note

Every TorchVision Dataset includes two arguments, transform and target_transform, that accept functions to modify the samples and the labels, respectively.

Let’s take a look at a preprocessed image. We can index a Dataset object like a list to retrieve an image and the corresponding label.

image, label = train_ds[0]

Image:
 - Type: <class 'torchvision.tv_tensors._image.Image'>
 - Min/Max: 0.0 - 1.0
 - Shape: 1 28 28

Label:
 - Type: <class 'int'>
 - Value: 6

Batching#

A Dataset retrieves the images and labels one sample at a time. While training a model, we typically want to pass samples in “batches” and reshuffle the data at every epoch. The DataLoader class abstracts this complexity for us in an easy API. It takes a Dataset object as an argument and provides a Python iterable over the dataset with support for automatic batching, multi-process data loading and many more features. We can configure a DataLoader with the following input arguments (look here for the full list).

  • batch_size: Number of samples to load per batch. Default is 1.

  • shuffle: If True, the data is reshuffled at every epoch. This is important for training.

  • num_workers: Number of subprocesses to use for data loading. The default, 0, means that the data will be loaded in the main process, which can slow down training for datasets where loading a sample takes a considerable amount of time (e.g., large images). For tiny datasets, 0 workers are usually faster.

  • drop_last: If True, the last batch is dropped in case it is smaller than the specified batch size. This occurs when the dataset size is not a multiple of the batch size. Only potentially helpful during training to keep a consistent batch size.

Here we define a batch size of 64 for the training set. As a result, each batch fetched by the dataloader will be a tensor of 64 images and a tensor of 64 labels.

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
valid_loader = DataLoader(valid_ds, batch_size=256, shuffle=False)
test_loader  = DataLoader(test_ds,  batch_size=256, shuffle=False)

Let’s retrieve the first batch of training samples.

for images, labels in train_loader:
    
    print("Image Batch Shape:", *images.shape)
    print("Label Batch Shape:", *labels.shape)

    break # without this, it will print for all batches

Image Batch Shape: 64 1 28 28
Label Batch Shape: 64

Note that the images and the labels are stacked along the first dimension (axis=0).

Important

In a tensor holding a batch of samples, the leading dimension (axis 0) conventionally corresponds to the sample index. Think of it as the axis that keeps track of the samples in a batch.

To better understand the data, we can visualize some of the images and labels in the first batch.

Hide code cell source

 
n = 10
side_by_side = torch.cat([img[0] for img in images[:n]], dim=1)

plt.figure(figsize=(10,5))
plt.imshow(side_by_side, cmap='gray')
plt.title("Labels: " + str(labels[:n].numpy()))
plt.axis('off')
plt.show()
../../_images/68dc9b49e9e6d5c868e1194eb628c2e6e426502e94b0036c5a8556a01741f873.png

Summary#

In this tutorial, we learned how to load a dataset, preprocess the data, and divide it into batches. These are the basic steps that are required to prepare data for training a neural network in PyTorch. In the next tutorial, we will learn how to define a neural network.