Linear probing vs fine-tuning

Setup#

pip install torch torchvision scikit-learn

Written and run against torch 2.13.0 (CPU), torchvision 0.28.0, scikit-learn 1.7.2.

This example uses a real dataset: fetch_olivetti_faces from scikit-learn. It holds 400 photographs of 40 people, 64 x 64 greyscale, about 4.5 MB. Plus the 44.7 MB ResNet-18 checkpoint. Total run time on a laptop CPU was under twenty seconds.

Both methods, same backbone, same split#

import time, torch, torch.nn as nn, torch.nn.functional as F, numpy as np
from sklearn.datasets import fetch_olivetti_faces
from sklearn.linear_model import LogisticRegression
from torch.utils.data import TensorDataset, DataLoader
from torchvision.models import resnet18, ResNet18_Weights

torch.manual_seed(0); np.random.seed(0)

faces = fetch_olivetti_faces()                    # 400 real photos of 40 people, ~4.5 MB
x = torch.tensor(faces.images).unsqueeze(1).repeat(1, 3, 1, 1)   # grey -> 3 channels
x = F.interpolate(x, size=112, mode="bilinear", align_corners=False)
mean = torch.tensor([0.485,0.456,0.406]).view(1,3,1,1)
std  = torch.tensor([0.229,0.224,0.225]).view(1,3,1,1)
x, y = (x - mean) / std, torch.tensor(faces.target)

tr = torch.cat([torch.arange(i*10, i*10+6) for i in range(40)])     # 6 per person to learn from
te = torch.cat([torch.arange(i*10+6, i*10+10) for i in range(40)])  # 4 per person to test on

def test_acc(model):
    model.eval()
    with torch.no_grad():
        pred = torch.cat([model(x[te][i:i+40]) for i in range(0, len(te), 40)]).argmax(1)
    return (pred == y[te]).float().mean().item()

# ---------- linear probe: freeze everything, fit one linear layer ----------
t0 = time.time()
backbone = resnet18(weights=ResNet18_Weights.DEFAULT)
backbone.fc = nn.Identity()                       # stop at the 512-number description
backbone.eval()
with torch.no_grad():
    feats = torch.cat([backbone(x[i:i+40]) for i in range(0, 400, 40)])
probe = LogisticRegression(max_iter=2000).fit(feats[tr], y[tr])
print(f"linear probe   accuracy {probe.score(feats[te], y[te]):.3f}   {time.time()-t0:.0f}s total")

# ---------- fine-tune: same backbone, every weight now moves ----------
t0 = time.time()
model = resnet18(weights=ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(512, 40)
opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
loader = DataLoader(TensorDataset(x[tr], y[tr]), batch_size=16, shuffle=True)
for ep in range(4):
    model.train()
    for xb, yb in loader:
        opt.zero_grad()
        F.cross_entropy(model(xb), yb).backward()
        opt.step()
    print(f"fine-tune ep{ep+1} accuracy {test_acc(model):.3f}   {time.time()-t0:.0f}s so far")

linear probe   accuracy 0.887   2s total
fine-tune ep1 accuracy 0.237   3s so far
fine-tune ep2 accuracy 0.613   5s so far
fine-tune ep3 accuracy 0.900   9s so far
fine-tune ep4 accuracy 0.969   13s so far

The accuracies repeated exactly across runs on one machine with the seeds set. The timings will not match yours — they depend on your CPU, thread count and background load. Repeated runs on the same laptop finished the fine-tune anywhere between 8 and 14 seconds. Read the ratio, not the numbers: the probe cost roughly a sixth of the fine-tune here.

What this run is telling you#

The probe hit 0.887 in two seconds, and never trained the network. Forty-way face identification, six examples per person. The backbone was trained on internet objects, never on faces. Random guessing here is 0.025.

Fine-tuning was worse than the probe for the first two epochs. 0.237, then 0.613. The head is random at the start. Its early gradients push the backbone around before carrying any useful signal. That is not a hyperparameter accident. It is the mechanism described by Kumar et al. in the researcher section.

Fine-tuning passed the probe at epoch three and finished at 0.969. Eight extra points of accuracy for roughly six times the compute, on 240 training images. Whether that trade is worth it is a product question, not a machine learning one.

The probe's real advantage is hidden in the code. Features were computed once, for all 400 images, then reused. Add ten more people and you extract features for the new photos only. Refit a logistic regression in a second, and ship. The fine-tuned model has to be retrained from the checkpoint.

Line by line#

backbone.fc = nn.Identity() replaces the classifier with a layer that returns its input unchanged. The model now outputs the 512-number description rather than class scores. This is cleaner than a forward hook and easier to read than slicing the model.

backbone.eval() before extraction is required, not optional. In train mode the BatchNorm layers rewrite their statistics as batches pass through. The "frozen" features would then depend on extraction order. See freezing and unfreezing layers.

LogisticRegression(max_iter=2000) is a linear classifier fitted by scikit-learn. It is the same maths as an nn.Linear head trained with cross-entropy. Its solver converges reliably on 240 rows, with no learning rate to pick. See logistic regression.

x[i:i+40] batching during extraction keeps peak memory low. 400 images at 112 x 112 through a ResNet-18 in one go is fine here. At 224 x 224 on a real dataset it is not.

Common mistakes#

Fitting the probe on features extracted in train mode. The features drift, the probe fits the drift, and the score does not reproduce. Fix: backbone.eval() and torch.no_grad().

Judging fine-tuning by epoch one. As the output above shows, epoch one can be far below the probe. Fix: run at least until the validation curve flattens before declaring a winner.

Comparing a tuned probe against an untuned fine-tune. Probing has one important knob, the regularisation strength C. Fine-tuning has learning rate, schedule, epochs, augmentation and freezing policy. An unfair comparison is the norm, in both directions. Fix: state the budget you gave each.

Extracting features once and then augmenting. Cached features cannot be augmented — the augmentation happens on pixels, upstream of the cache. Fix: extract several augmented copies per image, or accept that probing gives up augmentation.

Try it yourself#

Cut the training set from 6 photos per person to 2 by changing i*10+6 to i*10+2 in the tr line. Predict which method suffers more before you run it. The answer is in the next lesson but one, on few-shot classification.

What to learn next#

• Where ImageNet pretraining stops helping — what to do when the frozen features are wrong for your task.
• Few-shot image classification — the regime where probing wins by a wide margin.
• CLIP — the backbone whose frozen features made probing a headline result.