PythiloLesson 20 of 30 · progress kept in this browser

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.

This lesson uses torch, which Python in this browser does not include. Read it here, and run and check it in the Pythilo app on iPhone or iPad.

Next: Convolutions