Fine-Tuning

Once your new head is trained, you can do better by letting the top layers of the old network adapt to your data as well. But these weights took weeks to learn, and a careless learning rate would wreck them. So we need to be gentle.

• Let the top layers adapt
• Those weights took weeks to learn
• Be gentle

Different speeds for different layers

We keep most of the network frozen, and we unfreeze only layer four and the final layer. The key move is giving these parts different learning rates. The old layer gets a tiny rate, ten times smaller than the new head, so that it only gets small adjustments. Run it. Most of the weights stay frozen. Only a few million are training, and each group has its own speed. This is called discriminative learning rates.

from torch import nn, optim
from torchvision import models

model = models.resnet18(weights=None)
model.fc = nn.Linear(512, 10)

for name, p in model.named_parameters():
    p.requires_grad = name.startswith(("layer4", "fc"))

opt = optim.AdamW([
    {"params": model.layer4.parameters(), "lr": 1e-4},
    {"params": model.fc.parameters(), "lr": 1e-3},
])

used = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("Trainable weights:", used)
print("Frozen weights:   ", sum(p.numel() for p in model.parameters()) - used)
print("Learning rates:", [g["lr"] for g in opt.param_groups])

Trainable weights: 8398858
Frozen weights:    2782784
Learning rates: [0.0001, 0.001]

What to do with your data

Here's a simple guide. With a small dataset that's similar to the original, just train the new head. With a large similar dataset, fine tune the upper layers. With a small and different dataset, train on features from earlier layers, and use strong augmentation. With a large and different dataset, you can fine tune everything, or even start fresh.

Summary