Data Augmentation and Transfer Learning

You have 500 labeled photos of a defect class and a pretrained backbone expecting thousands. Augmentation is how you stretch those 500 images: apply label-preserving transforms so each epoch the model sees a different variation of every sample, which acts as regularization and breaks spurious shortcuts (background, watermark, camera angle) the model would otherwise memorize. The word doing the work is label-preserving — a horizontal flip of a screw thread is still the same defect; a horizontal flip of a chest X-ray where laterality matters is a different diagnosis.

The torchvision pipeline, split correctly between train and validation:

from torchvision import transforms

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),          # skip if chirality matters
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],   # ImageNet stats —
                         std=[0.229, 0.224, 0.225]),    # match the backbone
])

val_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

Three rules separate a pipeline that helps from one that silently poisons your metrics:

1. Validation transforms are deterministic. Resize + center crop, no randomness. If every eval run scores a different validation accuracy, you cannot compare experiments, and small improvements become noise. Randomness belongs to training only.
2. Split before you augment. Augmented variants of one source image must never straddle train and validation — a flipped crop of a validation photo sitting in training means you are scoring memorization. Split the raw file list first, then build the loaders.
3. Normalization must match the pretrained backbone. The frozen network learned features on ImageNet mean/std; feeding it pixel values scaled 0-1 while its weights expect normalized inputs shifts every activation, and transfer accuracy collapses by double-digit points for a reason that has nothing to do with your model.

Augmentation strength is a hyperparameter, not a virtue. Too weak and the model still memorizes backgrounds; too strong and labels start lying — random crops that crop out the defect, hue shifts that flip a ripeness class. Automated policies exist for exactly this reason: RandAugment applies N random operations at one shared magnitude M, so a small sweep over M replaces pages of hand-tuning. Watch the right signal while sweeping — heavier augmentation should widen the train loss gap while validation improves; if both worsen, the policy is destroying labels. Reproducibility is the last discipline: seed torch, numpy, and random, and log the full transform config with the run, because a transform change is a model change — two runs without their augmentation settings recorded are not comparable.

Which transforms are valid depends on your task, and interviewers like probing exactly there: horizontal flips are wrong for digits, letters, and any "left vs right" signal; rotations are wrong for text detection; for time series you cannot rotate or flip at all — jitter, scaling, and time warping are the analogous tools. Color jitter (hue especially) is wrong when color is the label, as in ripeness or rash classification.

A useful discipline: print what the pipeline actually emits. RandomResizedCrop at scale=(0.7, 1.0) can crop away the defect you are trying to detect — if 1 in 20 augmented images no longer contains the label, you have injected label noise at scale. Viewing a batch of transformed images before training costs two minutes and catches that. Once your data pipeline holds up, the next question is how to reuse a network trained on someone else's 1.2 million images. The architecture side — what the backbone is and why its head looks the way it does — came from CNN Architecture Patterns; here the concern is the training strategy: freeze what transfers, train what does not.

Transfer learning starts from weights someone else paid to train — torchvision's ResNet18_Weights.IMAGENET1K_V1 is 11.7M parameters of edge, texture, and shape detectors that transfer across almost any visual task. What changes between tasks is the last mapping from features to your labels, so the first strategy is brutally simple: freeze every pretrained weight and train only the new head.

import torch, torch.nn as nn
from torchvision import models

model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

for p in model.parameters():          # 1) freeze the backbone...
    p.requires_grad = False

model.fc = nn.Linear(512, 7)          # 2) ...then attach a NEW head, which
                                      #    starts trainable by default
print(sum(p.numel() for p in model.parameters() if p.requires_grad))  # 3,591

opt = torch.optim.Adam(model.fc.parameters(), lr=1e-3)   # 3) head params only

Order matters in that snippet: freeze before replacing the head. Had you frozen after, model.fc would be frozen too, and the optimizer would receive zero trainable parameters — a silent no-op where the loss never moves. Counting trainable parameters and confirming they match your head arithmetic (512 x 7 + 7 = 3,591) is the two-line sanity check for any freeze.

Two mechanics catch people in review:

• Frozen does not mean inert. Parameters with requires_grad=False skip the gradient graph, but the forward pass still runs, so a head-only run on a 10k-image set still costs full backbone compute. The cheap trick: run the frozen backbone once over the dataset, cache the 512-dim features to disk, then train on the cached matrix — head training becomes seconds per epoch.
• BatchNorm keeps updating unless you stop it. model.eval() would freeze BN running statistics — correct for inference — but you are training, and plain model.train() lets the backbone's BN layers recompute statistics from your 500 images, which corrupts features that were fine. The fix is to hold the frozen portion in eval mode while the head trains, and to keep the backbone's BN statistics fixed for the whole fine-tune unless your batches are large (64+) and your domain is far from ImageNet.

Audit the freeze in one line before the first epoch rather than discovering it in the loss curve:

print([(n, p.numel()) for n, p in model.named_parameters() if p.requires_grad])
# [('fc.weight', 3584), ('fc.bias', 7)] — nothing from layer1..layer4 appears

The cache path is more valuable than it sounds. Feature extraction is a pure forward job, so torch.inference_mode() around the loop halves memory and removes the autograd graph entirely; the resulting matrix (N x 512) trains a logistic regression head in seconds per epoch, which lets you iterate on regularization, class weights, and threshold choice before committing to any GPU time. The discipline that matters: whatever hyperparameters you pick on cached features are exactly the head configuration you then attach to the backbone for end-to-end runs — no silent differences between the probe stage and the real stage.

Is a head-only model enough? Read the training curve, not the vibes. If head training reaches 63% train / 61% val, the features are fine and the head underfits — move to unfreezing. If it reaches 99% train / 66% val, the head is memorizing 500 images and the features are not separating your classes — unfreezing with a high learning rate will make that worse, not better. Diagnosing which of those two you are in, slice by slice, is the discipline in Vision Model Error Analysis; before you get there, the middle option between "frozen" and "train everything" is fine-tuning, covered next.

Head-only training plateaus when your data looks unlike ImageNet — satellite tiles, spectrograms, microscopy — because 512 generic features cannot separate classes that need domain-specific ones. Fine-tuning unfreezes upper backbone layers and continues training them, but at a learning rate 10-100x lower than the head, and that ratio is the entire trick. Pretrained weights are already good; a learning rate sized for a randomly initialized head would overwrite them before the head itself converges. The failure has a name: catastrophic forgetting — 50 epochs later your features are trashed and validation is worse than the frozen baseline you never re-ran.

model.layer4.requires_grad_(True)          # last residual stage (512-dim)

param_groups = [
    {"params": model.layer4.parameters(), "lr": 1e-5},   # backbone: tiny steps
    {"params": model.fc.parameters(),     "lr": 1e-3},   # head: normal steps
]
opt = torch.optim.Adam(param_groups)
# keep BN statistics frozen even while layer4 trains:
for m in model.modules():
    if isinstance(m, nn.BatchNorm2d):
        m.eval()

These are called discriminative learning rates, and PyTorch supports them natively through parameter groups — one optimizer, per-group multipliers with multiplicative schedules if you prefer lr * 0.1 style. The staged recipe that works reliably on small datasets:

Gradual unfreezing means walking down that table — each stage starts from the weights the previous stage produced, so the model never takes a giant uncontrolled step away from good features. Stage 3 exists only if stage 2 still shows head underfitting; with a few hundred images per class, full fine-tuning routinely loses to stage 2 because the added capacity just fits noise, and no learning-rate setting fixes that.

Two mechanics make staged training safe in code:

def set_stage(model, trainable_prefixes):
    for name, p in model.named_parameters():
        p.requires_grad = name.startswith(trainable_prefixes)
    for m in model.modules():               # BN statistics stay frozen
        if isinstance(m, nn.BatchNorm2d):
            m.eval()

set_stage(model, "layer4")                  # stage 2
opt = torch.optim.Adam(filter(lambda p: p.requires_grad,
                              model.parameters()), lr=1e-5)

Rebuild the optimizer whenever the trainable set changes. Adam carries per-parameter moments; handing it the same parameter objects after freezing or unfreezing stages mixes stale statistics with fresh gradients, and a freshly constructed optimizer per stage avoids the question entirely. If overfitting still wins at stage 2, one more lever exists before you give up: L2-SP, a penalty toward the pretrained point p0 (loss + λ·Σ‖p - p0‖²) instead of toward zero — it anchors fine-tunes to features that already worked, which is exactly the inductive bias scarce data needs. Stop stages early with patience 3-5 on validation loss; a stage that has not moved in three epochs will not move in thirty.

The decision evidence lives in one comparison: always keep a frozen baseline checkpoint and a fine-tuned checkpoint evaluated on the same validation set. Fine-tuning that moves val accuracy 71% -> 74% while train sits at 99% told you the ceiling is features or labels, not capacity — which is exactly the fork the next topic takes apart: Vision Model Error Analysis is how you decide whether the next week goes to unfreezing more layers, collecting hard cases, or fixing labels. Log the learning-rate groups, the stage boundaries, and the frozen baseline's numbers with every run — a fine-tune without its probe comparison is an unverifiable claim.