03 · Transfer Learning 🔁

5 min read

Why train a vision model from scratch when someone already spent thousands of GPU-hours teaching one to see? Transfer learning lets you stand on their shoulders.

1. The 'Why'

Training a strong CNN from random weights needs two things most people don't have: a massive labeled dataset (think millions of images) and a lot of compute. But here's the liberating insight — the early layers of a CNN learn very general visual features: edges, corners, textures, color gradients. These features are useful for almost any vision task, whether you're classifying flowers, detecting defects on a production line, or reading X-rays. So instead of relearning "what an edge looks like" from scratch, you can grab a model already trained on a giant dataset like ImageNet and reuse its learned feature extractor. This is transfer learning, and it's arguably the most practical skill in applied computer vision.

The payoff is enormous: with transfer learning you can often build an excellent classifier from a few hundred images in minutes, where training from scratch would need hundreds of thousands of images and hours. The workflow is intuitive once you see it. You load a pretrained model, freeze its feature-extracting layers so their hard-won weights don't change, and replace its final classification head with a fresh one sized for your number of classes. Then you train only that new head (or, for fine-tuning, gently unfreeze some deeper layers too). torchvision ships dozens of pretrained models ready to go, making this a few lines of code — and the natural culmination of everything you've learned about models, layers, and training.

2. Core Concepts

Pretrained models. torchvision.models provides architectures (ResNet, EfficientNet, etc.) with weights already trained on ImageNet. Load them with weights=... to get those trained parameters.

Freezing. Setting param.requires_grad = False on the feature-extractor layers stops autograd from computing their gradients, so the optimizer leaves them untouched. This preserves the general features and speeds up training.

Replacing the head. The pretrained model's final layer outputs ImageNet's 1000 classes. You swap it for a new nn.Linear (or small stack) with out_features equal to your class count. The new layer starts random and is trainable.

Feature extraction vs. fine-tuning. Feature extraction: freeze everything except the new head. Fine-tuning: also unfreeze some later blocks and train them at a small learning rate to adapt features to your domain. Start with feature extraction; fine-tune if you need more.

Match the preprocessing. Your inputs must be normalized the same way the pretrained model expects (its documented mean/std and input size), or the features are garbage — connecting back to Module 04's transforms.

3. Code in Action

Loading a pretrained model and freezing it

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

# Load ResNet-18 with ImageNet-pretrained weights
weights = models.ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)

# Freeze every parameter so the feature extractor stays fixed
for param in model.parameters():
    param.requires_grad = False          # no gradients -> optimizer won't update these

Replacing the classification head

import torch.nn as nn

# resnet18's final layer is `model.fc`, a Linear(512 -> 1000 ImageNet classes)
num_features = model.fc.in_features      # 512 for resnet18
model.fc = nn.Linear(num_features, 10)   # NEW head for our 10 classes (trainable by default)

# Only the new head's parameters require grad, so only they will train:
trainable = [n for n, p in model.named_parameters() if p.requires_grad]
print(trainable)                         # ['fc.weight', 'fc.bias']

Training only the head (and the matching transform)

import torch
from torchvision import models

# Give the optimizer ONLY the parameters that require grad
optimizer = torch.optim.Adam(
    (p for p in model.parameters() if p.requires_grad),  # = the new head
    lr=1e-3,
)

# Use the EXACT preprocessing the pretrained weights were trained with
preprocess = weights.transforms()        # correct resize + normalize for this model
# train_ds = ImageFolderDataset(paths, labels, transform=preprocess)
# ...then the standard training loop from Module 05.

Optional: fine-tuning deeper layers

# After the head is trained, unfreeze the last block and continue at a SMALL lr
for param in model.layer4.parameters():
    param.requires_grad = True           # let the deepest features adapt

optimizer = torch.optim.Adam(
    (p for p in model.parameters() if p.requires_grad),
    lr=1e-4,                             # smaller lr to avoid wrecking good features
)

4. Common Pitfalls

Pitfall 1 — Mismatched preprocessing. A pretrained model expects inputs normalized and sized exactly as during its original training. Feed it differently-normalized images and the transferred features misfire, tanking accuracy. Always use the model's own weights.transforms() (or its documented mean/std and input size).

Pitfall 2 — Handing frozen parameters to the optimizer. If you pass model.parameters() wholesale after freezing, the optimizer includes frozen tensors (harmless but wasteful) — worse, if you forgot to freeze, you accidentally retrain the whole network. Pass only requires_grad=True parameters and verify which layers are trainable.

Pitfall 3 — Fine-tuning everything at a high learning rate. Unfreezing all layers and training at a large lr can destroy the pretrained features before the model adapts (catastrophic forgetting). Start with feature extraction; if you fine-tune, unfreeze gradually and use a much smaller learning rate.

5. Further Reading & Watch List

⬅️ Prev: 02 · CNN Architecture · ➡️ Next: 04 · Image Classification Project