All modulesS22 · CNN II, transfer learning▾
You rarely train an image network from scratch. You take a model already trained on millions of images (ImageNet) and adapt it to your problem. With little data, that's the difference between a good score and a miserable one. The model has already learned to see edges, textures and shapes; you just reorient it toward your classes.
Why transfer works
A CNN trained on many images learns general features in the lower layers: edges, corners, textures. These are useful for almost any image problem, not just the one it was trained on. The upper layers learn specific features (shapes of dogs, of cars). The transfer idea: you keep the general part and rewrite only the specific part.
That way, instead of training millions of weights from scratch, with millions of images you don't have, you use what the model already learned and train only a little, with little data. It's the most practical technique in all of image deep learning.
Feature extraction vs fine-tuning
There are two ways to adapt a pretrained model, depending on how much data you have.
- Feature extraction: you freeze the whole network (it no longer trains) and replace only the classification head, which you train on your classes. Fast, suited to when you have very little data.
- Fine-tuning: you also unfreeze the upper layers of the network and train them with a small learning rate, to nudge them toward your problem. Stronger, suited to when you have somewhat more data.
from torchvision import models
import torch.nn as nn
net = models.resnet18(weights="DEFAULT")
for p in net.parameters():
p.requires_grad = False # freeze everything (feature extraction)
net.fc = nn.Linear(net.fc.in_features, NUM_CLASSES) # new, trainable headNormalization: the same as at training
A pretrained model saw images normalized in a certain way: subtracted and divided by the ImageNet statistics (the means and deviations on the three channels). If you feed it images normalized differently, the input doesn't look like what it saw at training and it works badly, sometimes inexplicably badly.
A small learning rate for fine-tuning
In fine-tuning, the pretrained layers already have good weights. If you train them with a large learning rate, the big steps break exactly the valuable features you wanted to keep. Use a small learning rate for them, sometimes even smaller than for the new head. That way you nudge them finely, not destroy them.
- Pretrained models have learned general features reusable on other problems.
- Feature extraction (freeze everything, train only the head) is for very little data.
- Fine-tuning (unfreeze the upper layers, small learning rate) is for somewhat more data.
- Normalize the images with the same statistics as at pretraining (ImageNet).
- A large learning rate in fine-tuning breaks the good features already learned.