Transfer learning for vision

Labeled images are expensive to collect. Transfer learning avoids starting from random weights. You take a network trained on a large dataset such as ImageNet, replace its final layer with one for your classes, and train from there. The early layers have learned edges and textures that transfer to many tasks, so you need far less labeled data.

The recipe

1. Replace the head. Swap the 1000-way ImageNet classifier for a layer with K outputs, one per class of yours.
2. Choose how much to train. Either freeze the backbone and train only the head, or train everything with a smaller learning rate on the early layers than on the head.
3. Augment to match your domain. Crops, flips and color jitter help, as long as they keep the labels valid.

Sketch, not run here: this needs torch and torchvision installed.

import torchvision as tv
import torch

weights = tv.models.ResNet18_Weights.DEFAULT
model = tv.models.resnet18(weights=weights)
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

Then train with a low learning rate on the backbone and a higher one on model.fc.

Choosing how much to fine-tune

For a small dataset, start with a frozen backbone and a linear head. This is a linear probe. It trains quickly and gives a baseline, so you can measure how much fine-tuning adds. For a large dataset, or a domain far from ImageNet, full fine-tuning usually wins.

What transfers and what does not

Transfer can also fail on the input itself. Single-channel images fed to a three-channel network need a choice: repeat the channel, or retrain the first layer. Repeating the channel is a surprisingly good baseline.

Common mistakes

• Keeping the ImageNet head. Its 1000 outputs do not match your classes.
• Using ImageNet normalization on other data. Single-channel X-rays need their own mean and standard deviation.
• Mismatched preprocessing. A random crop in training with a center crop at serving shifts the input distribution. Use the same transform in both.

Check yourself

1. Which layer of a pretrained ResNet should you replace first, and why? The final classifier, because its outputs are for ImageNet's 1000 classes, not yours.
2. Why train a linear probe before full fine-tuning? It is a fast baseline, and the gap to full fine-tuning tells you whether the extra work is worth it.

Key takeaways

• Start from a pretrained backbone unless your domain is very different and your dataset is very large.
• Replace the head, then decide how much of the backbone to train.
• Early features transfer well. Normalization and the classifier often do not.
• Keep preprocessing identical between training and serving.

Further learning

Pooling and cheap convolutions shape the cost of these backbones: pooling and cheap convolutions.