Freezing and fine-tuning
Fine-tuning reuses a trained network for a new task: keep its body, which turns inputs into useful features, and train a new head on top. With little data, train only the head.
Freezing means requires_grad_(False) on every parameter of the body. Autograd then computes no gradients for them, which also saves time and memory, and you give the optimizer only the parameters that still train.
import torch
from torch import nn
body = nn.Sequential(nn.Linear(8, 16),
nn.ReLU())
for p in body.parameters():
p.requires_grad_(False)
head = nn.Linear(16, 3)
model = nn.Sequential(body, head)
trainable = [p for p in model.parameters()
if p.requires_grad]
print(sum(p.numel() for p in trainable))Freezing stops gradients, not every change: batch normalisation still updates its running averages in training mode. Call body.eval() after model.train() if the body has such layers.
Your task
Write freeze(module) that turns off gradients for every parameter of module and returns it. Write count_trainable(model) returning how many numbers in model still train. Write fine_tune(body, features, classes, x, y, steps=200): freeze body, return nn.Sequential(body, nn.Linear(features, classes)) trained with torch.optim.Adam at lr=0.05 on only the trainable parameters and nn.CrossEntropyLoss.
The explanation opens here once every check passes.