Transfer learning and fine-tuning

Torchvision and Hugging Face backbones, freezing parameters correctly, replacing heads, discriminative learning rates, and feature extraction as a baseline.

Loading a pretrained backbone

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

weights = ResNet50_Weights.IMAGENET1K_V2
model = models.resnet50(weights=weights)

# the transform that matches these exact weights
preprocess = weights.transforms()
print(preprocess)

# freeze the backbone, then replace the classifier head
for param in model.parameters():
    param.requires_grad = False

n_features = model.fc.in_features
model.fc = nn.Sequential(
    nn.Dropout(0.2),
    nn.Linear(n_features, 4),          # 4 output classes, logits
)

# only the new head has gradients
trainable = [p for p in model.parameters() if p.requires_grad]
print(sum(p.numel() for p in trainable))      # 2052 for 512->4

opt = torch.optim.AdamW(trainable, lr=1e-3, weight_decay=1e-4)

Use weights.transforms() rather than hand-writing normalisation. It is guaranteed to match the checkpoint, including crop size and interpolation.
Freezing sets requires_grad=False, which stops gradient computation. But batch-norm layers still update their running statistics unless you call model.eval() during the frozen stage.
Replacing the head changes the parameter count, so any saved optimiser state from a previous run no longer matches the model.
Freeze before you build the optimiser, or the frozen parameters end up in a param group with a learning rate and are silently skipped.

Progressive unfreezing and discriminative rates

def unfreeze_from(model, block_names):
    for name, child in model.named_children():
        if name in block_names:
            for param in child.parameters():
                param.requires_grad = True

def trainable_groups(model, base_lr=1e-4):
    """Give earlier layers a smaller rate than later ones."""
    groups, names = [], []
    for name, module in model.named_children():
        params = [p for p in module.parameters() if p.requires_grad]
        if params:
            groups.append(params)
            names.append(name)
    lrs = [base_lr * (1.5 ** i) for i in range(len(groups))]
    return [{"params": p, "lr": lr} for p, lr in zip(groups, lrs)], names

# stage 1: head only
opt1 = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-3)

# stage 2: unfreeze the last block, then recompile the optimiser
unfreeze_from(model, {"layer4"})
groups, names = trainable_groups(model, base_lr=1e-4)
print(dict(zip(names, [g["lr"] for g in groups])))
opt2 = torch.optim.AdamW(groups, weight_decay=1e-4)

Stage	Frozen	Learning rate	Why
1. Head only	Everything else	`1e-3`	Random head must not damage pretrained features
2. Last block	Early layers	`1e-4`	Adapt task-specific features gently
3. Full network	Nothing	`1e-5`	Only worthwhile with plenty of data
Feature extraction	Backbone, cached vectors	n/a	Train a linear model on frozen features

⚠️

When you change requires_grad, build a new optimiser. An optimiser keeps references to the parameter groups it was given at construction, so a parameter unfrozen afterwards is never updated — the loss plateaus and looks like an architecture problem.

Feature extraction as a fast baseline

import torch
from torch.utils.data import DataLoader

@torch.no_grad()
def extract_features(model, loader, device="cuda"):
    model.eval()
    feats, labels = [], []
    for images, targets in loader:
        images = images.to(device)
        out = model(images)                    # the backbone, head removed
        feats.append(out.flatten(1).cpu())
        labels.append(targets)
    return torch.cat(feats), torch.cat(labels)

trunk = torch.nn.Sequential(*list(model.children())[:-1])   # everything but fc
X_train, y_train = extract_features(trunk, train_loader)
X_val, y_val = extract_features(trunk, val_loader)

from sklearn.linear_model import LogisticRegression
clf = LogisticRegression(max_iter=1000).fit(X_train.numpy(), y_train.numpy())
print(clf.score(X_val.numpy(), y_val.numpy()))

# or a linear probe trained with PyTorch, which scales to larger data
probe = torch.nn.Linear(X_train.shape[1], 4)
opt = torch.optim.AdamW(probe.parameters(), lr=1e-3)
loss_fn = torch.nn.CrossEntropyLoss()

for epoch in range(30):
    for xb, yb in DataLoader(list(zip(X_train, y_train)), batch_size=256, shuffle=True):
        opt.zero_grad(set_to_none=True)
        loss = loss_fn(probe(xb), yb)
        loss.backward()
        opt.step()

Extract once, train many times. Caching frozen features turns a GPU-hour problem into a seconds-long one and makes the baseline essentially free.
A linear probe establishes how separable your classes are in the pretrained representation. If it scores well, full fine-tuning will help only a little.
Use a weighted sampler or class-weighted loss when the extracted feature set is imbalanced; the probe is otherwise dominated by the majority class.
Keep the augmentation that was used to produce the features fixed. Changing augmentation between feature extraction and probe training invalidates the comparison.

FAQ

How do I know whether to freeze or fine-tune?

Start frozen as a baseline, measure the validation score, then unfreeze the last block and measure again. If the second number is not clearly better, fine-tuning is not worth the compute on this dataset.

Why does the validation metric get worse after unfreezing?

Almost always the learning rate. Fine-tuning needs a rate one to two orders of magnitude smaller than head training, and unfreezing too many layers at once compounds the problem.

Modules, data and saving Learning-rate schedules and regularisation

Last refreshed 2026-09-18.

Transfer learning and fine-tuning

Loading a pretrained backbone

Progressive unfreezing and discriminative rates

Feature extraction as a fast baseline

FAQ

Related