TopolearnWith pretrained models, a few hundred labelled images can give results that would otherwise need millions.
Transfer learning reuses a network trained on a large dataset (such as ImageNet) as a feature extractor for your task. torchvision makes it a few lines:
weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
preprocess = weights.transforms()The weights object also provides the exact preprocessing (resize, crop, normalisation) the model was trained with; using different preprocessing quietly hurts accuracy.
The standard recipe:
1. Freeze the pretrained parameters with p.requires_grad_(False). 2. Replace the head: the final layer (model.fc for ResNets, model.classifier for many others) with a new nn.Linear sized for your classes. New layers are trainable by default. 3. Train only the head, passing just the trainable parameters to the optimizer. 4. Optionally fine-tune: unfreeze some or all layers and continue with a learning rate 10-100x smaller, often via parameter groups.
Frozen BatchNorm layers still update their running statistics in train() mode; for very small datasets, keep the backbone in eval() mode.
The same pattern applies to text and audio models from the Hugging Face hub, which are PyTorch nn.Modules under the hood.
Example
import torch
from torch import nn
from torchvision.models import resnet18, ResNet18_Weights
weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights) # downloads ImageNet weights once
preprocess = weights.transforms() # matching resize + normalise
for p in model.parameters():
p.requires_grad_(False) # freeze the backbone
model.fc = nn.Linear(model.fc.in_features, 5) # new 5-class head, trainable
trainable = [p for p in model.parameters() if p.requires_grad]
print(sum(p.numel() for p in trainable)) # 2565 (512 * 5 + 5)
opt = torch.optim.AdamW(trainable, lr=1e-3)
x = preprocess(torch.rand(3, 300, 400)).unsqueeze(0) # fake image
print(model(x).shape) # torch.Size([1, 5])Remember
- Use the weights' own
transforms()for preprocessing - Freeze, replace the head, train the head, then fine-tune gently
- Give the optimizer only trainable parameters
Watch out for
- Fine-tuning the whole network with a large learning rate, destroying the pretrained features.