PyTorch - Datasets

PyTorch - Working with Datasets and DataLoaders

Efficient data handling is crucial for training machine learning models. PyTorch provides two main primitives for this: torch.utils.data.Dataset and torch.utils.data.DataLoader. These utilities help in abstracting the process of accessing and iterating over your data.

The torchvision library, a companion to PyTorch for computer vision tasks, includes several popular datasets out-of-the-box within torchvision.datasets.

Core Concepts: Dataset and DataLoader

torch.utils.data.Dataset

An abstract class representing a dataset. Any custom dataset you create should inherit from Dataset and override two methods:

• __len__(self): Should return the total number of samples in the dataset.
• __getitem__(self, index): Should return the sample (e.g., an image and its label) at the given index.

torchvision.datasets provides pre-built Dataset classes for many common datasets.

torch.utils.data.DataLoader

An iterator that provides features like:

• Batching: Grouping multiple samples into a batch.
• Shuffling: Randomly shuffling the data at every epoch to prevent model bias.
• Parallel Loading: Using multiple subprocesses to load data in parallel, speeding up training.

DataLoader takes a Dataset object as input and makes it easy to iterate over the data.

Transforms

When working with datasets, especially in computer vision, you often need to apply transformations to your data (e.g., resizing images, converting them to tensors, normalizing pixel values). torchvision.transforms provides common image transformations.

Two types of transforms can be applied when loading datasets:

• transform: A function/transform that takes in a sample (e.g., a PIL Image) and returns a transformed version. Common transforms include transforms.ToTensor(), transforms.Normalize(), transforms.Resize().
• target_transform: A function/transform that takes in the target (e.g., a label) and returns a transformed version.

These can be composed together using transforms.Compose().

Example: Using torchvision.datasets.MNIST

MNIST is a classic dataset of handwritten digits. Let’s see how to load it.

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# Define a transform to convert images to tensors and normalize them
# Normalization values for MNIST are typically mean=0.1307, std=0.3081
transform_mnist = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Download and load the training data
# root: directory where data is/will be stored
# train=True: get the training set
# download=True: download if not already present
# transform: apply the defined transformations
train_dataset_mnist = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    transform=transform_mnist,
    download=True
)

# Download and load the test data
test_dataset_mnist = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    transform=transform_mnist,
    download=True
)

# Create DataLoaders
# batch_size: number of samples per batch
# shuffle=True: shuffle training data to ensure randomness
train_loader_mnist = DataLoader(dataset=train_dataset_mnist, batch_size=64, shuffle=True)
test_loader_mnist = DataLoader(dataset=test_dataset_mnist, batch_size=1000, shuffle=False)

print(f"Number of training samples: {len(train_dataset_mnist)}")
print(f"Number of test samples: {len(test_dataset_mnist)}")

# Example of iterating through the DataLoader
data_iter = iter(train_loader_mnist)
images, labels = next(data_iter)
print(f"Shape of a batch of images: {images.shape}") # Expected: torch.Size([64, 1, 28, 28])
print(f"Shape of a batch of labels: {labels.shape}") # Expected: torch.Size([64])

Parameters for torchvision.datasets classes like MNIST:

• root (str): Root directory where the dataset MNIST/processed/training.pt and MNIST/processed/test.pt exist or will be saved.
• train (bool, optional): If True, creates dataset from training.pt, otherwise from test.pt.
• download (bool, optional): If True, downloads the dataset from the internet and puts it in root directory. If dataset is already downloaded, it is not downloaded again.
• transform (callable, optional): A function/transform that takes in a PIL image and returns a transformed version.
• target_transform (callable, optional): A function/transform that takes in the target and transforms it.

Example: Using torchvision.datasets.CIFAR10

CIFAR10 is another popular dataset with 10 classes of 32x32 color images.

# Define transforms for CIFAR10 (3-channel images)
# Normalization values for CIFAR10 are typically mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]
transform_cifar10 = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

train_dataset_cifar10 = torchvision.datasets.CIFAR10(
    root='./data',
    train=True,
    transform=transform_cifar10,
    download=True
)

train_loader_cifar10 = DataLoader(dataset=train_dataset_cifar10, batch_size=64, shuffle=True)

print(f"\nNumber of CIFAR10 training samples: {len(train_dataset_cifar10)}")

data_iter_cifar = iter(train_loader_cifar10)
images_cifar, labels_cifar = next(data_iter_cifar)
print(f"Shape of a CIFAR10 batch of images: {images_cifar.shape}") # Expected: torch.Size([64, 3, 32, 32])

Other Notable torchvision.datasets

torchvision offers many other datasets, including:

• ImageNet: A large-scale image dataset (requires manual download due to size and licensing).
• COCO (Common Objects in Context): For object detection, segmentation, and captioning. Example:

# COCO Example (requires COCO API and data to be downloaded separately)
# try:
#   coco_cap = torchvision.datasets.CocoCaptions(
#       root = 'path/to/coco/images/train2017', # Directory with images
#       annFile = 'path/to/coco/annotations/captions_train2017.json', # Annotation JSON file
#       transform=transforms.ToTensor()
#   )
#   print(f'Number of COCO caption samples: {len(coco_cap)}')
#   if len(coco_cap) > 0:
#       img, target = coco_cap[0] # Get first sample (image and its captions)
#       print(f'COCO Image shape: {img.shape}, Number of captions: {len(target)}')
# except Exception as e:
#   print(f"Could not load COCO dataset. Ensure data and API are set up. Error: {e}")

The output for COCO might look like (if data is available):

Number of COCO caption samples: 118287 (Example number for train2017 captions)
COCO Image shape: torch.Size([3, height, width]), Number of captions: 5 (Example number)

Creating Custom Datasets

For most real-world projects, you’ll need to create your own custom Dataset class. This involves:

1. Creating a class that inherits from torch.utils.data.Dataset.
2. Implementing __init__(self, ...): Initialize your dataset, e.g., load file paths, annotations.
3. Implementing __len__(self): Return the total number of samples.
4. Implementing __getitem__(self, idx): Load and return one sample (e.g., image and label) for the given index idx. This is where you typically apply transforms.

For more detailed information, refer to the official PyTorch documentation on Dataset and DataLoader and explore tutorials on creating custom datasets.

Additional Resources:

• PyTorch Data Loading Tutorial: https://pytorch.org/tutorials/beginner/basics/data_tutorial.html