Transfer learning

Two ways into the same idea

Before You Start

The Core Idea

Transfer borrows the eyesight of a model that saw millions of images. Here you meet the model's two parts, freeze one and replace the other, measure the gain honestly, then choose between three strategies.

1Backbone and head

A trained model has two parts. The backbone is convolutional layers turning an image into a feature vector, and the head is a final linear layer turning features into classes.

Look at the picture: in ResNet18 the backbone is nearly all the weights. It is the general part that knows how to see, and the head is specific to the original task.

2Freezing and swapping the head

Freezing means stopping the backbone layers from training with requires_grad = False, so their weights stay fixed. Training takes minutes, not hours, even on a CPU.

Then we put in a new head with our class count. In ResNet it is called model.fc, and in other models model.classifier.

This is the real shape of transfer learning. It downloads ResNet weights trained on ImageNet, so it is here for reading.

from torch import nn
from torchvision.models import ResNet18_Weights, resnet18

model = resnet18(weights=ResNet18_Weights.DEFAULT)
for p in model.parameters():
    p.requires_grad = False                         # freeze the backbone
model.fc = nn.Linear(model.fc.in_features, 4)       # new head for 4 classes
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

Freeze everything, then replace the fc head with a new layer for your class count. The new layer is created after freezing, so it is trainable automatically, and the optimiser gets only its parameters. 512 features by 4 classes: about two thousand parameters train, out of 11 million.

3Measure across seeds

How much transfer helps depends on what the original model saw. So we run each method with five seeds and compare the mean and the range.

Look at the picture: the three bars are very close. The gap between them is smaller than the swing between seeds.

Without internet we cannot download ImageNet weights here, so we build a small "pretraining" ourselves: a model learns digits 0 to 4 with all their data, then we transfer it to digits 5 to 9 with only 25 images. Three methods, five seeds each.

import statistics

import torch
from sklearn.datasets import load_digits
from torch import nn

digits = load_digits()
X = torch.tensor(digits.images, dtype=torch.float32).unsqueeze(1) / 16
y = torch.tensor(digits.target)
order = torch.randperm(len(y), generator=torch.Generator().manual_seed(0))
X, y = X[order], y[order]

source = y < 5                                   # "pretraining" task: digits 0-4, all labels
X_src, y_src = X[source], y[source]
X_tgt, y_tgt = X[~source], y[~source] - 5        # new task: digits 5-9
X_few, y_few = X_tgt[:25], y_tgt[:25]            # only 25 labelled images
X_test, y_test = X_tgt[300:], y_tgt[300:]


def backbone():
    return nn.Sequential(
        nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        nn.Flatten(),
    )


def train(model, X, y, params):
    opt = torch.optim.Adam(params, lr=0.01)
    for _ in range(150):
        loss = nn.functional.cross_entropy(model(X), y)
        opt.zero_grad()
        loss.backward()
        opt.step()


def accuracy(model):
    with torch.no_grad():
        return (model(X_test).argmax(1) == y_test).float().mean().item()


results = {"scratch": [], "frozen": [], "fine-tuned": []}
for seed in range(5):
    torch.manual_seed(seed)
    pretrained = backbone()
    train(nn.Sequential(pretrained, nn.Linear(128, 5)), X_src, y_src, pretrained.parameters())

    scratch = nn.Sequential(backbone(), nn.Linear(128, 5))
    train(scratch, X_few, y_few, scratch.parameters())
    results["scratch"].append(accuracy(scratch))

    frozen_body = backbone()
    frozen_body.load_state_dict(pretrained.state_dict())
    frozen_body.requires_grad_(False)
    frozen = nn.Sequential(frozen_body, nn.Linear(128, 5))
    train(frozen, X_few, y_few, frozen[1].parameters())
    results["frozen"].append(accuracy(frozen))

    tuned_body = backbone()
    tuned_body.load_state_dict(pretrained.state_dict())
    tuned = nn.Sequential(tuned_body, nn.Linear(128, 5))
    train(tuned, X_few, y_few, tuned.parameters())
    results["fine-tuned"].append(accuracy(tuned))

for name, accs in results.items():
    print(f"{name:10} mean {statistics.mean(accs):.3f}  min {min(accs):.3f}  max {max(accs):.3f}")

The honest result: the gaps are small and fall within the seed range. Frozen is slightly higher on average and more stable (worst seed 87.6% versus 85.1%), and full fine-tuning on 25 images is slightly worse because it damages the features and memorises. The reason is that the "pretraining" here is only 900 small images, and 8 by 8 digits are easy enough for a small model to learn from scratch. With ImageNet (1.2 million varied images) and a real photo task, the gap is usually tens of percentage points. The practical lesson: compare across seeds, and never assume transfer always wins.

4Three strategies and graded rates

Always start by freezing the whole backbone, because it is the baseline. With somewhat enough data unfreeze the last block, and with plenty of data or very different data, like X-rays, fine-tune the whole model.

Look at the picture: the more data you have, the further along you go. And when you fine-tune everything, each layer gets its own learning rate.

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

The closer a layer is to the input, the smaller its rate. Early layers know edges, which are much the same in every image, so you do not want to change them much. Later layers are more specific to the original task and need more adjustment.

How It Works, Step by Step: how to transfer a trained model to your classes

Worked example

A trained model has two parts: a backbone (the convolutional layers) turning an image into a feature vector, and a head (a final linear layer) turning features into classes. The head is specific to the original task; the backbone is general.

Try it without code

Myth vs Reality

Try it yourself

Quick review cards

Check Your Understanding