Data Pipelines and Augmentation for Vision Models

The classifier in the last post started overfitting at epoch 10. The usual first response is augmentation — but a badly chosen augmentation makes things worse, and some quietly destroy your labels. This post covers which ones help, which ones lie, and how to keep the GPU fed.

Prerequisites: Your first image classifier.

What augmentation actually claims

The definition contains the trap. Every augmentation is a claim: this change does not alter the correct answer. When the claim is true, you get free training data. When it is false, you are training the model on wrong labels.

Look at the hue shift row. In the recycling classifier, the glass–plastic confusion was the dominant error, and the little colour tint each material gives is one of the few real cues. Shifting hue randomly deletes that cue and makes the worst confusion worse.

Watch a crop go wrong

RandomResizedCrop is the single most effective augmentation for classification and the easiest to misconfigure. Its scale parameter is the fraction of the original area kept.

PyTorch’s default is scale=(0.08, 1.0). On a 500×500 photo, 8% of the area is $0.08 \times 250{,}000 = 20{,}000$ px², a region of about 141×141 pixels, which is then upscaled to 224×224.

The third panel is a training example labelled “plastic” that contains no plastic. Feed the model enough of those and you teach it that empty background predicts plastic.

The default works for ImageNet because with 1.2 million images, occasional garbage crops are diluted. With 840 training images they are a meaningful fraction of every epoch.

Measuring what each augmentation is worth

Do not guess. Add one at a time and measure against a fixed validation set.

Two things to read off that chart. Training accuracy should fall as you augment; if it stays at 99.8% your augmentation is doing nothing. And the hue-shift dip is real and reproducible, because it removes the tint cue the glass–plastic decision depends on.

Mixup and CutMix

These are stranger than they look, and they work.

With $\lambda = 0.7$, image A being cardboard and image B being glass:

$\tilde{x} = 0.7 \, x_A + 0.3 \, x_B$ $\tilde{y} = 0.7 \, [1,0,0,0] + 0.3 \, [0,0,1,0] = [0.7,\; 0,\; 0.3,\; 0]$

Trace one pixel. If A has red 200 and B has red 60 at the same position:

$0.7(200) + 0.3(60) = 140 + 18 = 158$

Check the second column: $0.7(180) + 0.3(80) = 126 + 24 = 150$. ✓

CutMix does the same with a rectangle instead of a blend: paste a patch of B into A, and set $\lambda$ to the fraction of area that stayed A. If the pasted patch covers 30% of the image, the label is 70% A and 30% B — the same arithmetic, but the result looks like a real image rather than a ghostly double exposure.

import numpy as np, torch

def mixup(x, y, alpha=0.4):
    lam = np.random.beta(alpha, alpha)
    idx = torch.randperm(x.size(0), device=x.device)
    return lam * x + (1 - lam) * x[idx], y, y[idx], lam

# in the training loop
xm, ya, yb, lam = mixup(x, y)
out = model(xm)
loss = lam * criterion(out, ya) + (1 - lam) * criterion(out, yb)

The last two rows matter for small projects. Mixup and CutMix are strong regularisers that need extra epochs to pay off, and on a 1,200-image dataset the extra training time is usually better spent collecting more images.

Keeping the GPU fed

An augmented pipeline is CPU work: decode the JPEG, resize, transform, normalise. Do too little of it in parallel and the GPU sits waiting.

Measure the per-image CPU cost:

One worker delivers 133 img/s. The GPU can consume ResNet-18 forward and backward at roughly 380 img/s. So:

train_dl = DataLoader(
    train_ds,
    batch_size=32,
    shuffle=True,
    num_workers=8,
    pin_memory=True,           # faster host-to-GPU copies
    persistent_workers=True,   # don't respawn workers every epoch
    prefetch_factor=4,         # batches queued per worker
    drop_last=True,            # keeps batch shapes constant
)

persistent_workers=True is the cheapest win here. Without it, PyTorch tears down and restarts every worker process at the end of each epoch, which on a small dataset can cost more time than the epoch itself.

The visualization above contrasts a correct pipeline with a leaky one. The distinction it draws — fit preprocessing on training data only — is exactly the trap in the next section.

Statistics must come from training data only

If you compute your own normalization statistics instead of using ImageNet’s, compute them on the training split alone.

# Correct: training split only
mean, std, n = 0.0, 0.0, 0
for x, _ in DataLoader(train_ds_no_aug, batch_size=64):
    b = x.size(0)
    x = x.view(b, x.size(1), -1)
    mean += x.mean(2).sum(0)
    std  += x.std(2).sum(0)
    n    += b
mean, std = mean / n, std / n

Computing them over the full dataset leaks information about the test set into training. The effect is small for normalization statistics specifically, but the habit is what matters: anything fitted to data must be fitted to training data only, and that applies to class weights, PCA, and any learned resize or crop policy too.

Test-time augmentation

At inference, run several augmented versions of each image and average the predictions.

@torch.no_grad()
def predict_tta(model, img):
    views = [img, torch.flip(img, dims=[3])]           # original + horizontal flip
    probs = [torch.softmax(model(v), dim=1) for v in views]
    return torch.stack(probs).mean(0)

Flip-only TTA is the sweet spot: 0.7 points for double the cost, which is fine for a batch job and usually fine for a phone app. Beyond that you pay 5× more compute for another 0.7 points.

Only use TTA transforms that are valid for your task. If horizontal flip would change the label — reading text, arrows, or gauges — flip TTA actively hurts.

graph TD
  A["Raw images on disk"] --> B["Split by session<br/>train / val / test"]
  B --> C["Compute stats on TRAIN only"]
  C --> D["Train: random augmentation"]
  C --> E["Val/Test: deterministic resize + crop"]
  D --> F["DataLoader<br/>workers tuned to saturate GPU"]
  E --> G["DataLoader<br/>shuffle=False"]
  F --> H["Training"]
  G --> H
  H --> I["Optional TTA at inference"]

Practice task

Use the classifier from the previous post, or any small dataset.

1. Train with no augmentation. Record train and validation accuracy and the gap.
2. Add RandomResizedCrop(224, scale=(0.7, 1.0)). Record the gap.
3. Add horizontal flip. Then colour jitter. Then rotation. One at a time.
4. Plot train and validation accuracy across all five runs.
5. Now set scale=(0.08, 1.0) and save 20 augmented images to disk. Look at them. How many contain no object?
6. Time one epoch at num_workers of 0, 2, 4, 8, and 16. Plot epoch time against workers.
7. Add flip TTA at inference and measure the accuracy change.
8. Add an augmentation you believe is wrong for your task — hue shift, or vertical flip if orientation matters. Measure how much it costs.

Step 5 is the one that changes behaviour permanently. Seeing your own training images cropped down to blank background, still carrying a confident class label, makes the scale parameter stop being an abstraction.

Summary

Every augmentation asserts that a change does not alter the label. Horizontal flip and mild crops are usually safe; hue shift cost 2.3 points on the recycling dataset because material tint was one of the few real cues separating glass from plastic.

RandomResizedCrop’s default scale=(0.08, 1.0) keeps as little as 8% of the area — about 141×141 pixels from a 500×500 photo — often producing a blank crop that still carries a class label. Use (0.5, 1.0) or (0.7, 1.0) on small datasets.

Add augmentations one at a time and watch the train–validation gap, which fell from 21.5 to 2.5 points here. Training accuracy falling is the intended effect, not a problem.

Mixup blends images and labels by the same $\lambda$: at $\lambda = 0.7$, pixel values 200 and 60 become 158, and the label becomes $[0.7, 0, 0.3, 0]$. Both mixup and CutMix need larger datasets and more epochs to pay off.

Finally, size the DataLoader from measurement. At 7.5 ms of CPU work per image, one worker supplies 133 img/s against a GPU that consumes 380, so three workers saturate it and sixteen only cost RAM.

What comes next

You have now trained and regularised a model while treating ResNet-18 as a black box. CNN architectures from LeNet to ResNet opens it up: what a convolution layer actually computes, why depth stalled at around 20 layers before residual connections, and how to read an architecture diagram well enough to choose between backbones on purpose.