Module 9 — Transfer learning with torchvision

Our Fashion-MNIST MLP reaches about 91 % accuracy. A convolutional network trained from scratch on the same data would reach 93 or 94 %. A pretrained ResNet18, adapted to Fashion-MNIST through transfer learning, gets past 95 % in one epoch. That gap — an epoch of fine-tuning versus days of training from scratch — is what makes transfer learning the default first move on any vision problem where labelled data is finite. This module shows exactly how to do it in PyTorch, and where the traps hide.

Why early layers transfer

A deep vision network learns a hierarchy of representations. The first convolutions detect edges, corners and colour gradients — features that describe local structure in any image, regardless of subject. Middle layers detect textures and small motifs. Late layers detect object parts specific to the training classes. The classification head is entirely specific to the original taxonomy.

Fashion-MNIST is grayscale clothing at 28-by-28. ImageNet is colour photographs of a thousand object categories at hundreds of pixels per side. Different domains, different resolutions — but the first two levels of the hierarchy are essentially identical. An edge detector remains an edge detector whether the pixels came from a leopard or a t-shirt. Keeping the early ResNet weights and retraining only the top is the whole idea of transfer learning.

Loading a pretrained model

torchvision.models ships every classical architecture with pretrained weights, indexed by dataset.

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

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

The old pretrained=True string is deprecated; the new API asks you to specify which set of weights explicitly, because torchvision provides several for many architectures — an original set and a refreshed one with more modern training recipes.

The weights object also carries the preprocessing transform the network expects:

preprocess = weights.transforms()
# Resize(256), CenterCrop(224), ToTensor, Normalize(imagenet_mean, imagenet_std)

Feeding raw Fashion-MNIST pixels to this network would fail: the model expects 3-channel colour images at 224-by-224, normalised with ImageNet statistics. The Fashion-MNIST inputs are 1-channel grayscale at 28-by-28. Two options bridge the gap: adapt the inputs (resize, replicate the channel three times, then apply the ImageNet normalisation), or adapt the network's first convolution (change it to accept one channel and drop the ImageNet normalisation). This course uses the first option, which is simpler and preserves the pretrained weights exactly.

from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

Freezing the backbone

Freezing means turning off gradient computation for a parameter, so that the optimiser cannot update it:

for param in model.parameters():
    param.requires_grad = False

This freezes every weight in the ResNet. The optimiser will still walk parameters() but every one of them will have requires_grad=False, and optimizer.step() will skip them. It is efficient — no gradient allocation for frozen parameters, no backward traversal past them — and it is what makes phase 1 of transfer learning cheap.

Replacing the head

ResNet18's classification head is a single nn.Linear layer named fc, mapping the 512-dim feature vector to 1000 ImageNet classes. Replacing it is one line:

model.fc = nn.Linear(model.fc.in_features, 10)

The new layer is randomly initialised, so requires_grad=True by default, which is exactly what we want. The rest of the network remains frozen.

model.fc.in_features is preferable to hardcoding 512: it lets the same code work on ResNet34 (512), ResNet50 (2048) or any other variant without adjustment.

Two-phase training

The effective procedure runs in two phases, and the order is not negotiable.

Phase 1: feature extraction. Backbone frozen, only the new head trains. Standard learning rate.

optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3, weight_decay=1e-4,
)
for epoch in range(3):
    train_one_epoch(...)

The filter restricts the optimiser to trainable parameters. Adam and AdamW ignore frozen ones through requires_grad, but keeping frozen parameters in the optimiser wastes memory on their moment estimates.

Phase 2: fine-tuning. Unfreeze the top of the backbone. Learning rate an order of magnitude lower.

for param in model.layer4.parameters():
    param.requires_grad = True

optimizer = torch.optim.AdamW(
    [{"params": model.layer4.parameters(), "lr": 1e-4},
     {"params": model.fc.parameters(),     "lr": 1e-3}],
    weight_decay=1e-4,
)
for epoch in range(5):
    train_one_epoch(...)

Two parameter groups with discriminative learning rates. The head, still adjusting to a new task, keeps a normal rate. The unfrozen backbone layer, close to a good solution already, uses a rate an order of magnitude smaller so gradients from the head do not scramble it.

Why phase 1 must precede phase 2

The new head is randomly initialised. On the very first batch, its output is meaningless, the loss is high, and gradients flowing back through the entire network are enormous. Applied to a pretrained ResNet at a normal learning rate, they erase in a few batches the representation accumulated over millions of ImageNet images. Phase 1 exists precisely to bring the head to a reasonable state before the backbone is exposed to gradients.

Skipping phase 1 and going straight to fine-tuning is not merely suboptimal: it is the mechanism by which most beginner transfer attempts fail. Validation accuracy climbs slowly, plateaus below the frozen-backbone baseline, and often regresses. The fix is not more capacity or more data; it is the two phases in order.

The BatchNorm ambiguity

ResNet contains many BatchNorm2d layers. Their running statistics were fit on ImageNet, and they matter as much as the convolution weights. Setting requires_grad=False freezes the learnable affine parameters (weight and bias) but does not stop the layers from updating their running statistics when the model is in train mode.

On a small Fashion-MNIST training set, those statistics drift toward what small clothing batches look like, away from what the frozen weights expect. Accuracy plateaus and can even regress — the exact symptom of the module 8 transfer-learning trap in TensorFlow, and it exists here too.

Two remedies. The safe one is to keep the backbone in eval mode during phase 1:

model.train()
model.layer4.eval()      # keep BN running stats frozen during phase 2
model.bn1.eval()         # and every other BN outside the trainable region

The more portable one, on modern torchvision, uses freeze_batchnorm utilities or an explicit walk over nn.BatchNorm2d submodules calling .eval() on each frozen one.

requires_grad=False freezes weights, not running statistics

This is the single most common transfer-learning mistake in PyTorch. Two separate mechanisms — requires_grad for gradients, .train() / .eval() for layer behaviour — must both be handled explicitly. On a base with dozens of BatchNorm layers, forgetting the second cancels most of the first's benefit.

How much to unfreeze

Situation	Strategy
Little data, similar domain	freeze all, train the head alone
Plenty of data, similar domain	unfreeze the last block (layer4)
Little data, distant domain	unfreeze middle layers (layer2, layer3)
Plenty of data, very distant domain	consider training from scratch

Fashion-MNIST is small and visually distant from ImageNet: grayscale versus colour, clothing versus everything, 28-by-28 versus 224-by-224. The reasonable strategy is to unfreeze layer4 for phase 2 and leave layer1 through layer3 frozen; that gets to 95 % validation accuracy without much tuning.

Sanity-check with a frozen-backbone baseline

Before unfreezing anything, run phase 1 to convergence and note the accuracy. Any fine-tuning strategy that ends up below that baseline is worse than not fine-tuning at all, and the failure mode is almost always the BatchNorm trap or a learning rate that is too high. The baseline is your safety net.

In summary

• Early layers transfer because they encode features — edges, textures — that describe local image structure independent of the original task or domain.
• Freezing means requires_grad=False; replace the head with a fresh nn.Linear sized to your classes, using in_features rather than a hardcoded 512.
• Train in two phases: feature extraction with the backbone frozen, then fine-tuning with the last block unfrozen at a learning rate one order of magnitude lower.
• requires_grad=False does not freeze BatchNorm running statistics; putting the frozen sub-modules in .eval() mode is the second, mandatory step.

Next module: TorchScript, ONNX and serving — how the fine-tuned ResNet finally leaves Python and answers HTTP requests.