Data Augmentation

Data Augmentation#

Data augmentation is a technique used to artificially expand a training dataset by generating plausible variations of existing data. For example, images can be augmented by applying rotations, flips, or color adjustments. This approach is a form of regularization that helps reduce overfitting by introducing more diversity into the training set. It is especially valuable when the available dataset is limited or imbalanced.

In this tutorial, we demonstrate how to apply data augmentation to cats-and-dogs images. Only the training set is augmented, while the validation and test sets are kept deterministic to ensure fair evaluation. This allows us to monitor the model’s performance accurately and reserve the test set for a final assessment.

Hide code cell source

 
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Subset
import torchvision
import torchvision.transforms.v2 as v2
from torchvision.datasets import ImageFolder

from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import pathlib

Preparation#

The previous tutorial downloaded and validated the cats-vs-dogs images. We reuse those images and recreate the same seeded train/validation/test split.

Hide code cell source

 
IMAGE_SIZE = 128
SPLIT_SEED = 42

# Deterministic preprocessing
preprocess = v2.Compose([
    v2.ToImage(),
    v2.Resize(IMAGE_SIZE, antialias=True),
    v2.CenterCrop(IMAGE_SIZE),
    v2.ToDtype(torch.float32, scale=True),
])

# Full dataset
data_path = pathlib.Path(".data") / "cats_vs_dogs" / "PetImages"
dataset = ImageFolder(data_path, transform=preprocess)

# Training split
all_idx = list(range(len(dataset)))
train_idx, holdout_idx = train_test_split(
    all_idx,
    stratify=dataset.targets,
    test_size=0.30,
    random_state=SPLIT_SEED,
)

# Validation and test splits
holdout_targets = [dataset.targets[i] for i in holdout_idx]
valid_idx, test_idx = train_test_split(
    holdout_idx,
    stratify=holdout_targets,
    test_size=0.50,
    random_state=SPLIT_SEED,
)

train_ds = Subset(dataset, train_idx)
valid_ds = Subset(dataset, valid_idx)
test_ds  = Subset(dataset, test_idx)

print(f"Train: {len(train_ds)}, validation: {len(valid_ds)}, test: {len(test_ds)}")

Train: 17497, validation: 3750, test: 3750

We define a simple convolutional network for the cats-vs-dogs dataset. We copied the model definition from the previous tutorial and paste it here for convenience.

class BaselineModel(torch.nn.Module):
    
    def __init__(self, image_size: int):
        super().__init__()
        ksize = 3
        self.conv1 = torch.nn.Conv2d(3, 32, ksize)
        self.conv2 = torch.nn.Conv2d(32, 64, ksize)
        self.conv3 = torch.nn.Conv2d(64, 128, ksize)
        self.conv4 = torch.nn.Conv2d(128, 128, ksize)
        flat_dim = self.__calc_dim(image_size)
        self.fc1 = torch.nn.Linear(flat_dim, 256)
        self.fc2 = torch.nn.Linear(256, 1)

    def forward(self, x):
        for conv in [self.conv1, self.conv2, self.conv3, self.conv4]:
            x = F.relu(conv(x))
            x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
    
    def __calc_dim(self, input_dim: int):
        """Returns the tensor size after the flatten layer."""
        n = input_dim
        for conv in [self.conv1, self.conv2, self.conv3, self.conv4]:
            n = n - conv.kernel_size[0] + 1 # CONV output size
            n = n // 2                      # POOL output size
        dim = n * n * self.conv4.out_channels
        return dim

Transforms#

TorchVision provides a list of transforms that can be used to preprocess and augment images for classification, detection, segmentation, and related tasks. The following are some of the most important points to keep in mind when using transforms.

  • Most transforms behave like torch.nn.Module objects and can be composed.

  • They commonly accept PIL images or channel-first tensors; many also accept batched tensors.

  • The value range of pixels is implicitly defined by the data type.

    • Tensor images with a float data type are expected to have values in [0, 1].

    • Tensor images with a uint8 data type are expected to have values in [0, 255].

    • Tensor images with an integer data type are expected to have values in [0, MAX_VALUE].

  • Resizing and geometric transforms involve interpolation; they have an antialias argument that can be used to reduce artifacts when downsampling.

We already encountered several transforms in the previous tutorials, such as ToImage and ToDtype. The first one converts a PIL image into TorchVision’s image tensor, whereas the second one converts the data type and rescales integer values consistently.

See also

This example illustrates the effect of various transforms provided by TorchVision.

Augmentation pipeline#

Our training pipeline uses the following transforms for data augmentation.

The figure below illustrates the augmentation pipeline. Note that the horizontal flip and the color jitter are skipped with a certain probability, based on a decision taken independently for each image in the batch. When a probability argument is not supported by a transform, which is the case for ColorJitter, we use RandomApply to apply a transform with a given probability.

Augmentation Pipeline

In addition to the augmentation transforms, the training pipeline also includes data conversion and rescaling. This ensure that the returned tensors contain floating-point values in the range [0, 1], which is the expected input range for the neural network.

augmentation = v2.Compose([
    v2.ToImage(),
    v2.RandomResizedCrop(IMAGE_SIZE, scale=(0.7, 1.0), antialias=True),
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomApply([v2.ColorJitter(0.2, 0.2, 0.2, 0.1)], p=0.2),
    v2.ToDtype(torch.float32, scale=True),
])

Augmented training set#

Previously, we created a training set without data augmentation. We will now create a second ImageFolder dataset over the same files that includes data augmentation. Only the training data will use augmentation. Validation and test data continue to use the deterministic pipeline defined earlier.

# Load the full dataset with augmentation
augmented_dataset = ImageFolder(data_path, transform=augmentation)

# Create the train set using the same indices as before
augmented_train_ds = Subset(augmented_dataset, train_idx)

Important

Note that ImageFolder loads images lazily, so augmentation runs each time a sample is requested rather than creating files on disk. This is online data augmentation. Repeated access usually produces a different view, although randomness does not guarantee that every result is unique.

Visualizing the augmented images#

Let’s visualize some images from the augmented training set to see the effect of the augmentation pipeline.

Hide code cell source

 
plt.figure(figsize=(8, 4), tight_layout=True)

for i in range(4):
    image, _ = train_ds[i]
    augmented, _ = augmented_train_ds[i]

    plt.subplot(2, 4, 1 + 2*i)
    plt.imshow(image.permute(1,2,0))
    plt.axis("off")
    plt.subplot(2, 4, 1 + 2*i + 1)
    plt.imshow(augmented.permute(1,2,0))
    plt.axis("off")
../../_images/adb4a0f4c58d150b4194ee53110b89ece6e97ff4becdcbff05d642f97b6ec0bc.png

Training with data augmentation#

We now train the model on the augmented training set, monitor the untouched validation set, and evaluate the test set only after the training is complete. We will use the same trainer as before. The only difference is that the model is trained on the augmented training set instead of the original one.

from training import Trainer
from torcheval.metrics import BinaryAccuracy

Training uses a different optimizer and epoch count than the previous tutorial. The results are thus not directly comparable. The goal here is to demonstrate how to use data augmentation in practice.

Warning: this cell is computationally expensive without a GPU.

def binary_adapter(model, batch, func):
    inputs, targets = batch
    logits = model(inputs).squeeze(-1)
    return func(logits, targets.float())

train_loader = DataLoader(augmented_train_ds, batch_size=64, shuffle=True)
valid_loader = DataLoader(valid_ds, batch_size=128, shuffle=False)
test_loader = DataLoader(test_ds, batch_size=128, shuffle=False)

torch.manual_seed(42)

device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model   = BaselineModel(IMAGE_SIZE).to(device)
loss_fn = torch.nn.BCEWithLogitsLoss()

optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, amsgrad=True)
epochs = 20

trainer = Trainer()
trainer.set_adapter(binary_adapter)

history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)

===== Training on cuda:0 device =====

Epoch  1/20: 100%|██████████| 274/274 [01:10<00:00,  3.89it/s, train_loss=0.6701, valid_loss=0.6469]
Epoch  2/20: 100%|██████████| 274/274 [01:16<00:00,  3.59it/s, train_loss=0.6234, valid_loss=0.5964]
Epoch  3/20: 100%|██████████| 274/274 [01:16<00:00,  3.58it/s, train_loss=0.5851, valid_loss=0.5540]
Epoch  4/20: 100%|██████████| 274/274 [01:14<00:00,  3.70it/s, train_loss=0.5526, valid_loss=0.5307]
Epoch  5/20: 100%|██████████| 274/274 [01:08<00:00,  4.00it/s, train_loss=0.5211, valid_loss=0.5080]
Epoch  6/20: 100%|██████████| 274/274 [01:06<00:00,  4.09it/s, train_loss=0.5016, valid_loss=0.5417]
Epoch  7/20: 100%|██████████| 274/274 [01:10<00:00,  3.90it/s, train_loss=0.4851, valid_loss=0.4750]
Epoch  8/20: 100%|██████████| 274/274 [01:09<00:00,  3.96it/s, train_loss=0.4588, valid_loss=0.4652]
Epoch  9/20: 100%|██████████| 274/274 [01:08<00:00,  4.02it/s, train_loss=0.4448, valid_loss=0.4356]
Epoch 10/20: 100%|██████████| 274/274 [01:07<00:00,  4.06it/s, train_loss=0.4314, valid_loss=0.4525]
Epoch 11/20: 100%|██████████| 274/274 [01:08<00:00,  4.01it/s, train_loss=0.4244, valid_loss=0.4162]
Epoch 12/20: 100%|██████████| 274/274 [01:08<00:00,  3.99it/s, train_loss=0.4113, valid_loss=0.4176]
Epoch 13/20: 100%|██████████| 274/274 [01:06<00:00,  4.12it/s, train_loss=0.3996, valid_loss=0.4022]
Epoch 14/20: 100%|██████████| 274/274 [01:07<00:00,  4.04it/s, train_loss=0.3955, valid_loss=0.3955]
Epoch 15/20: 100%|██████████| 274/274 [01:07<00:00,  4.05it/s, train_loss=0.3886, valid_loss=0.3816]
Epoch 16/20: 100%|██████████| 274/274 [01:07<00:00,  4.04it/s, train_loss=0.3830, valid_loss=0.3839]
Epoch 17/20: 100%|██████████| 274/274 [01:07<00:00,  4.04it/s, train_loss=0.3782, valid_loss=0.3720]
Epoch 18/20: 100%|██████████| 274/274 [01:07<00:00,  4.07it/s, train_loss=0.3700, valid_loss=0.3692]
Epoch 19/20: 100%|██████████| 274/274 [01:08<00:00,  4.00it/s, train_loss=0.3607, valid_loss=0.3649]
Epoch 20/20: 100%|██████████| 274/274 [01:08<00:00,  3.98it/s, train_loss=0.3572, valid_loss=0.3901]

We plot the training loss along with the validation loss. Because augmented training examples are intentionally harder and change between epochs, training loss may exceed validation loss.

Hide code cell source

 
plt.plot(history['train_loss'], label='Train loss')
plt.plot(history['valid_loss'], label='Validation loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()
../../_images/44541c8e6b58ff7bbf32d55d36ea16676a434d0116db06bb9ca2a08a947849fc.png

After the training decisions are fixed, we evaluate the model accuracy on the test set.

Hide code cell source

 
trainer.set_metrics(accuracy=BinaryAccuracy())
ans = trainer.eval(model, test_loader)

print(f"Test Accuracy: {ans['accuracy']:.2%}")

Test Accuracy: 84.77%

Summary#

In this notebook, we learned how to use data augmentation in PyTorch. We created an augmentation pipeline, applied it to a subset of the cats-vs-dogs dataset, and trained a simple ConvNet on the augmented data.

Important

Data augmentation applies random, label-preserving transformations to training samples on the fly. Validation and test samples are not augmented, since it would make the evaluation inconsistent.

Data augmentation is a powerful technique to increase the diversity of the training data and improve the generalization of the model. However, it is important to be careful with the choice of transformations, as some of them may not be suitable for the problem at hand. For example, flipping an image horizontally may not be a good idea for a dataset of handwritten digits.

Note

Specialized libraries like Albumentations provide a wide range of data augmentation techniques and can be easily integrated into PyTorch pipelines.

The following table shows the data augmentation techniques used in some popular ConvNets for image classification, object detection, and segmentation tasks.

Model

Data Augmentations

LeNet-5

Translate, Scale, Squeeze, Shear

AlexNet

Translate, Flip, Intensity Changing

ResNet

Crop, Flip

DenseNet

Flip, Crop, Translate

MobileNet

Crop, Elastic distortion

NasNet

Cutout, Crop, Flip

ResNeSt

AutoAugment, Mixup, Crop

DeiT

AutoAugment, RandAugment, Random erasing, Mixup, CutMix

Swin Transformer

RandAugment, Mixup, CutMix, Random erasing

U-Net

Translate, Rotate, Gray value variation, Elastic deformation

Faster R-CNN

Flip

YOLO

Scale, Translate, Color space

SSD

Crop, Resize, Flip, Color Space, Distortion

YOLOv4

Mosaic, Distortion, Scale, Color space, Crop, Flip, Rotate, Random erase, Cutout, Hide and Seek, GridMask, Mixup, CutMix, StyleGAN