Fine-tuning a pretrained network, in code

Almost nobody trains a vision or text model from scratch. You start from one that has already learned general features and adapt it with the few thousand examples you have. This section does that in code, in the order that works: a linear probe on frozen features first, as a baseline that fine-tuning has to beat and often does not beat by much, then a new head, then unfreezing the later layers at a much lower rate. That order also avoids the standard way of wrecking a pretrained network, which is letting a randomly initialised head push large gradients back through it on the first step.

Start with the cheapest version#

Freeze the pretrained network, run your images through it once to get feature vectors, and train a logistic regression on them. This linear probe takes minutes and is a surprisingly strong baseline. Fine-tuning must beat it to be worth the effort.

The word once is the point. Because nothing in the backbone changes, the features are the same on every epoch, so you compute them a single time and then train on a small matrix of numbers instead of on images. That is why it costs seconds rather than hours, and it is also the fastest way to try twenty different classifiers.

import torch, torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights
from sklearn.linear_model import LogisticRegression

device = "cuda" if torch.cuda.is_available() else "cpu"

backbone = resnet18(weights=ResNet18_Weights.DEFAULT)
backbone.fc = nn.Identity()                       # keep the features, drop the classifier
backbone = backbone.to(device).eval()

@torch.no_grad()                                  # no graph, no gradients, half the memory
def extract(loader):
    F, Y = [], []
    for x, y in loader:
        F.append(backbone(x.to(device)).cpu()); Y.append(y)
    return torch.cat(F).numpy(), torch.cat(Y).numpy()

Ftr, ytr = extract(train_loader)                  # [n, 512] once, then reused
Fte, yte = extract(test_loader)
print(LogisticRegression(max_iter=3000).fit(Ftr, ytr).score(Fte, yte))
# 0.8040 on 2,000 CIFAR-10 images, in 21 seconds

Replace the head and train#

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

model = resnet50(weights=ResNet50_Weights.DEFAULT)
for p in model.parameters():
    p.requires_grad = False                         # freeze everything

model.fc = nn.Linear(model.fc.in_features, 5)       # new head for 5 classes, trainable

opt = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)

The order of those lines matters. Freezing happens first and the new head is attached afterwards, because a freshly created Linear has requires_grad True by default and would otherwise be frozen along with everything else. On a ResNet-50 head for 5 classes that leaves a fraction of a per cent of the network trainable.

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(trainable, total)      # 5,130 of 11,181,642 on a ResNet-18 with 10 classes = 0.046%

for p in m.parameters():
    p.requires_grad = False

before = m.bn1.running_mean.clone()
m.train(); m(torch.randn(4, 3, 224, 224))
print(torch.allclose(before, m.bn1.running_mean))    # False  -> it changed

m.eval();  before = m.bn1.running_mean.clone()
m(torch.randn(4, 3, 224, 224))
print(torch.allclose(before, m.bn1.running_mean))    # True   -> now it is really frozen

Then unfreeze, gently#

Once the new head has trained, unfreeze the later layers and continue with a much lower learning rate. Early layers detect general things like edges and textures and need little change; later layers are more task-specific. Giving earlier layers a smaller learning rate respects that.

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

opt = torch.optim.AdamW([
    {"params": model.layer4.parameters(), "lr": 1e-5},
    {"params": model.fc.parameters(),     "lr": 1e-4},
])

A list of dicts is how PyTorch expresses per-group hyperparameters, and lr is not the only one that can vary; weight_decay works the same way, which matters because decaying the pretrained weights towards zero is a way of erasing exactly what you came for.

Per-layer learning rates across the whole network#

Staged unfreezing needs you to decide how many stages and when to move between them. The alternative is to unfreeze everything at once and let the learning rate encode the same intent continuously: smallest at the input, rising by a fixed multiplier per block, largest at the head. This is discriminative fine-tuning, and on a small dataset it usually matches staged unfreezing in less wall-clock time.

groups = [m.conv1, m.bn1, m.layer1, m.layer2, m.layer3, m.layer4, m.fc]
base, decay = 1e-3, 2.6                # 2.6 is the fast.ai default, and as good as any

params = [
    {"params": g.parameters(), "lr": base / decay ** (len(groups) - 1 - i)}
    for i, g in enumerate(groups)
]
opt = torch.optim.AdamW(params)

stage     learning rate    parameters     how much it is allowed to move
conv1       3.2e-06             9,408        essentially nothing
bn1         8.4e-06               128
layer1      2.2e-05           147,968
layer2      5.7e-05           525,568
layer3      1.5e-04         2,099,712
layer4      3.8e-04         8,393,728        the bulk of the network, and most of the change
fc          1.0e-03             5,130        random at step 0, so it needs the full rate

conv1 moves 310 times more slowly than the head.

The head still has a problem this does not solve: it is random, so its first gradients are large and they travel backwards into layers that are now unfrozen. Two standard fixes, and either is enough. Train the head alone for one epoch first, then apply the per-layer rates. Or use a warmup, so the first few hundred steps happen at a fraction of the rate and the head has stopped being random by the time full-size updates arrive.

Use the right preprocessing#

A pretrained model expects inputs prepared the way they were during its training: size, normalisation, tokenisation. The weights object in torchvision, and the tokenizer that ships with a Hugging Face model, provide exactly that. Using your own preprocessing quietly costs accuracy.

weights = ResNet18_Weights.DEFAULT
transform = weights.transforms()          # never write these numbers yourself
print(transform)

# ImageClassification(
#     crop_size=[224]
#     resize_size=[256]
#     mean=[0.485, 0.456, 0.406]          # ImageNet channel means
#     std=[0.229, 0.224, 0.225]
#     interpolation=InterpolationMode.BILINEAR
# )

Getting this wrong is the most common silent loss of accuracy in transfer learning, because nothing errors: the model receives images with the wrong mean and scale, its first layer responds to values it never saw during pretraining, and every feature downstream is slightly off. You lose a few points and blame the dataset. Use the model's own transform for validation, and build your augmentation around its normalisation rather than instead of it.

Forgetting#

Fine-tuning hard on a small dataset can destroy general abilities the model had, known as catastrophic forgetting. Signs are a large jump in training accuracy with poor validation, or failure on inputs slightly different from your training set. Lower learning rates, fewer unfrozen layers, fewer epochs and methods that update only small added weights, such as LoRA, reduce it.

The mechanism is not mysterious. Gradient descent minimises the loss on the data in front of it and has no term for anything else, so weights encoding abilities your dataset never exercises are free to be overwritten in service of the current objective. Nothing in the loss protects them, and the model gives no indication that it has stopped being able to do something you never asked it about.

How to notice                                  What it looks like
hold out data from BEFORE fine-tuning          accuracy on the original task, measured each epoch
evaluate on inputs slightly off-distribution   a cliff rather than a gentle decline
compare against the frozen probe every epoch   fine-tuning drops below its own starting point
track the distance from the original weights   sum of squared differences, growing fast

Defence                        Cost                      How much it helps
lower learning rate            slower convergence         most of the benefit, for free
freeze more layers             less adaptation possible   a lot, and it is one line
fewer epochs, early stopping   may underfit               a lot; forgetting grows with steps
LoRA or an adapter             a little extra code        the base weights never change at all
mix in original-task data      you need that data         the direct fix, when available
weight decay towards the
  ORIGINAL weights, not zero   a custom optimiser step    effective, and rarely implemented

LoRA deserves the emphasis it gets, because it changes the shape of the problem rather than trading against it. The pretrained weights are frozen and a small pair of low-rank matrices is trained alongside them, so the original model is arithmetically untouched and the adaptation is a separate few megabytes you can remove, swap, or serve several of against one base model. That last property is why it is the default for language models, where keeping one copy of a large model and many small adapters is the difference between one deployment and fifty.