Unit 07.00: Feature extraction vs fine-tuning

Someone has already trained a model on far more images than you have. There are two ways to use that, and they suit different situations.

Freeze it, or continue training it

Feature extraction freezes the pretrained backbone and trains only a new head. Fast, needs little data, and cannot damage what the backbone already knows.

Fine-tuning continues training the backbone too. More capacity to adapt, but it needs more data and a much smaller learning rate or it destroys the pretrained weights.

import torch
from torch import nn

torch.manual_seed(0)
# Stand in for a pretrained backbone. In practice this comes from torchvision.
backbone = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(),
                         nn.AdaptiveAvgPool2d(1), nn.Flatten())
images = torch.randn(6, 3, 32, 32)

# Feature extraction: freeze the backbone, train only a new head.
for p in backbone.parameters():
    p.requires_grad = False
head = nn.Linear(16, 2)

trainable = sum(p.numel() for p in backbone.parameters() if p.requires_grad) \
    + sum(p.numel() for p in head.parameters())
total = sum(p.numel() for p in backbone.parameters()) + sum(p.numel() for p in head.parameters())
print(f"feature extraction: {trainable} trainable of {total} total")

# Fine-tuning: unfreeze the backbone too.
for p in backbone.parameters():
    p.requires_grad = True
trainable_ft = sum(p.numel() for p in backbone.parameters() if p.requires_grad) \
    + sum(p.numel() for p in head.parameters())
print(f"fine-tuning       : {trainable_ft} trainable of {total} total")

# Rule of thumb: little data or a similar domain -> extract features.
# Plenty of data or a distant domain -> fine-tune, at a small learning rate.

Compare the trainable parameter counts: feature extraction trains a tiny fraction of the total, fine-tuning trains all of it.

The rule of thumb: little data, or a domain close to what the backbone saw, favours extraction. Plenty of data, or a distant domain, favours fine-tuning.

The mistake this prevents

Fine-tuning on a few hundred images at the default learning rate. The pretrained weights are overwritten in the first few steps and you end up with a randomly initialised model that took longer to train.

Takeaway

Start frozen. Unfreeze only if the frozen version is not good enough, and lower the learning rate when you do.

Mark complete