How to use data augmentation pytorch

PytorchHow-ToBeginner · 4 min read

How to Use Data Augmentation in PyTorch for Better Models

Use torchvision.transforms to apply data augmentation in PyTorch by composing transformations like RandomHorizontalFlip, RandomRotation, and ColorJitter. Apply these transforms to your dataset during loading to create varied training data that helps your model generalize better.

📐

Syntax

Data augmentation in PyTorch is done using torchvision.transforms. You create a transforms.Compose object that chains multiple transformations. Each transform modifies the input image randomly or deterministically.

transforms.RandomHorizontalFlip(): flips image horizontally with 50% chance.
transforms.RandomRotation(degrees): rotates image randomly within given degrees.
transforms.ColorJitter(): randomly changes brightness, contrast, saturation, and hue.
transforms.ToTensor(): converts PIL image to PyTorch tensor.

Apply this composed transform when loading your dataset.

python

import torchvision.transforms as transforms

data_transforms = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor()
])

💻

Example

This example shows how to apply data augmentation to the CIFAR10 training dataset using torchvision.datasets and DataLoader. It demonstrates random flips, rotations, and color changes on images during training.

python

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define data augmentation transforms
transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.1, contrast=0.1),
    transforms.ToTensor()
])

# Load CIFAR10 training dataset with augmentation
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# Iterate one batch and print shape
images, labels = next(iter(train_loader))
print(f'Batch image tensor shape: {images.shape}')
print(f'Batch labels tensor shape: {labels.shape}')

Output

Files already downloaded and verified Batch image tensor shape: torch.Size([64, 3, 32, 32]) Batch labels tensor shape: torch.Size([64])

⚠️

Common Pitfalls

Common mistakes when using data augmentation in PyTorch include:

Applying augmentation to validation or test data, which should remain unchanged for fair evaluation.
Not converting images to tensors after augmentation, causing errors in model input.
Using augmentation transforms that distort labels (e.g., random crop without adjusting bounding boxes in object detection).
Forgetting to normalize images after augmentation, which can affect model training.

python

import torchvision.transforms as transforms

# Wrong: Applying augmentation to test data
transform_test_wrong = transforms.Compose([
    transforms.RandomHorizontalFlip(),  # Should NOT be here for test
    transforms.ToTensor()
])

# Right: Only convert to tensor and normalize for test
transform_test_right = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

📊

Quick Reference

Here is a quick summary of common data augmentation transforms in PyTorch:

Transform	Description
RandomHorizontalFlip()	Flip image horizontally randomly
RandomVerticalFlip()	Flip image vertically randomly
RandomRotation(degrees)	Rotate image randomly within degrees
ColorJitter()	Randomly change brightness, contrast, saturation, hue
RandomCrop(size)	Crop random region of given size
ToTensor()	Convert PIL image to PyTorch tensor
Normalize(mean, std)	Normalize tensor image with mean and std

✅

Key Takeaways

Use torchvision.transforms.Compose to chain multiple data augmentation steps.

Apply augmentation only to training data, not validation or test sets.

Always convert images to tensors after augmentation for model input.

Common augmentations include random flips, rotations, color jitter, and crops.

Normalize images after augmentation to help model training stability.