Unit 07.01: Freezing and unfreezing layers

Freezing has two halves, and doing only the first is a bug that produces no error at all.

requires_grad, and the optimiser's parameter list

Setting requires_grad = False stops gradients being computed for a tensor. But the optimiser holds its own list of parameters, and if a frozen tensor is in that list it can still be updated by momentum or weight decay from earlier steps.

Do both: set the flag, and construct the optimiser from only the trainable parameters.

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 8), nn.ReLU(), nn.Linear(8, 1))

# Freeze everything except the last layer.
for p in model.parameters():
    p.requires_grad = False
for p in model[-1].parameters():
    p.requires_grad = True

opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=0.05)
before = [p.clone() for p in model.parameters()]

x, y = torch.randn(32, 8), torch.randn(32, 1)
for _ in range(20):
    opt.zero_grad()
    nn.MSELoss()(model(x), y).backward()
    opt.step()

for i, (old, new) in enumerate(zip(before, model.parameters())):
    changed = not torch.allclose(old, new)
    print(f"parameter {i}: changed = {changed}")

# Only the last layer moved. Note the second half of the rule: the optimiser
# must also be given only the trainable parameters, or it will happily update
# tensors you thought were frozen once you unfreeze them later.

The output confirms only the last layer's parameters changed. Everything upstream is byte-identical after twenty optimisation steps.

That check — snapshot the parameters, train, compare — is the reliable way to verify freezing worked. Reading the code is not.

The mistake this prevents

Freezing by setting requires_grad = False but passing model.parameters() to the optimiser anyway. It appears to work, and then behaves strangely when you later unfreeze.

Takeaway

Freeze with the flag *and* the optimiser's parameter list. Verify by comparing weights before and after.

Mark complete