Unit 06.02: Freezing, unfreezing and the learning rate

Fine-tuning has two phases and two learning rates, and using one rate for both destroys the pretrained weights.

Head first, then unfreeze at a much lower rate

Freezing the backbone and training only the classifier, then unfreezing with a far smaller learning rate.

The code freezes a small model and reports the parameter split.

import torch
from torch import nn

model = nn.Sequential(
    nn.Conv2d(3, 8, 3, padding=1), nn.ReLU(),
    nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(),
    nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(16, 2))

for param in list(model.parameters())[:-2]:
    param.requires_grad = False

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad)
print(f"frozen parameters   : {frozen:,}")
print(f"trainable parameters: {trainable:,}  ({trainable / (trainable + frozen):.1%})")

print("\ntypical recipe:")
for phase, lr, what in [("1. head only", "1e-3", "backbone frozen"),
                        ("2. fine-tune", "1e-5", "unfreeze, much lower LR")]:
    print(f"  {phase:16} lr {lr:6} {what}")

# The second learning rate is the part people get wrong. Unfreezing at the
# original rate destroys the pretrained weights in a few batches -- which is
# the whole thing you were trying to keep.

Almost all the parameters are frozen in phase one, so the classifier learns against a fixed feature extractor. That is fast and stable.

The second learning rate is the part people get wrong. Unfreezing at the original rate pushes large gradients through weights that took enormous compute to learn, and destroys them in a few batches - which shows up as a model that trains and performs worse than the frozen version.

The mistake this prevents

The mistake is unfreezing everything immediately to "let the model adapt". With a randomly initialised head, the first gradients are large and meaningless, and they flow straight into the pretrained layers.

Takeaway

Train the head with the backbone frozen, then unfreeze at a much lower learning rate. Unfreezing early or at the original rate destroys what you were transferring.

Mark complete