Machine Learning
Challenge 4
Write a training loop.
Challenge 4: freeze, then fine-tune
In the final challenge you'll apply the two stages of transfer learning. You'll prepare the "pretrained" body yourself by training it on a side task, then, for a new task, train first only the head and then the whole network at a small rate:
# 1. preparation: body + head_A, on task A
# 2. stage 1: body frozen, head_B trained
# 3. stage 2: body open (small lr), body + head_B together
Checking the trainable parameter count at every stage verifies that the freezing really took effect.
Exercise Preparation and stage 1
Train task A (the moon data) for 300 steps with a body (Linear(2, 32) → ReLU → Linear(32, 32) → ReLU) and baş_A = Linear(32, 2). Then freeze the body; for task B (the moon with flipped labels: 1 - y), train a new baş_B = Linear(32, 2) for 100 steps. Print the stage-1 trainable parameter count and that task B's accuracy is above 0.95 on one line.
Show example solution
import math
import torch
from torch import nn
torch.manual_seed(0)
def ay(n):
t = torch.rand(n) * math.pi
ust = torch.stack([torch.cos(t), torch.sin(t)], 1)
alt = torch.stack([1 - torch.cos(t), 0.5 - torch.sin(t)], 1)
X = torch.cat([ust, alt]) + 0.1 * torch.randn(2 * n, 2)
y = torch.cat([torch.zeros(n), torch.ones(n)]).long()
return X, y
X, y = ay(200)
govde = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU())
bas_A = nn.Linear(32, 2)
kayip_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(list(govde.parameters()) + list(bas_A.parameters()), lr=0.01)
for _ in range(300):
opt.zero_grad()
kayip_fn(bas_A(govde(X)), y).backward()
opt.step()
for p in govde.parameters():
p.requires_grad = False
bas_B = nn.Linear(32, 2)
model = nn.Sequential(govde, bas_B)
opt = torch.optim.Adam(bas_B.parameters(), lr=0.01)
for _ in range(100):
opt.zero_grad()
kayip_fn(model(X), 1 - y).backward()
opt.step()
egitilebilir = sum(p.numel() for p in model.parameters() if p.requires_grad)
with torch.no_grad():
print(egitilebilir, (model(X).argmax(1) == 1 - y).float().mean().item() > 0.95)
Exercise Stage 2: fine-tuning
After stage 1 (the code is given), open the body and set up Adam with parameter groups of lr=1e-4 for the body and lr=1e-3 for the head. Fine-tune for 100 steps. Print the trainable parameter count and that task B's accuracy is still above 0.95 on one line.
Show example solution
import math
import torch
from torch import nn
torch.manual_seed(0)
def ay(n):
t = torch.rand(n) * math.pi
ust = torch.stack([torch.cos(t), torch.sin(t)], 1)
alt = torch.stack([1 - torch.cos(t), 0.5 - torch.sin(t)], 1)
X = torch.cat([ust, alt]) + 0.1 * torch.randn(2 * n, 2)
y = torch.cat([torch.zeros(n), torch.ones(n)]).long()
return X, y
X, y = ay(200)
govde = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU())
bas_B = nn.Linear(32, 2)
model = nn.Sequential(govde, bas_B)
kayip_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
for _ in range(400):
opt.zero_grad()
kayip_fn(model(X), 1 - y).backward()
opt.step()
for p in govde.parameters():
p.requires_grad = False
for p in govde.parameters():
p.requires_grad = True
opt = torch.optim.Adam([{"params": govde.parameters(), "lr": 1e-4}, {"params": bas_B.parameters(), "lr": 1e-3}])
for _ in range(100):
opt.zero_grad()
kayip_fn(model(X), 1 - y).backward()
opt.step()
egitilebilir = sum(p.numel() for p in model.parameters() if p.requires_grad)
with torch.no_grad():
print(egitilebilir, (model(X).argmax(1) == 1 - y).float().mean().item() > 0.95)
Test yourself
-
Which parameters should be given to the optimiser in stage 1?
- Only the new head's parameters
- All parameters
- Only the body
- None
Show answer
Only the new head's parameters — Since the body is frozen it won't be updated; only the head is given to the optimiser.
-
Make the body trainable again.
for p in govde.parameters():
p.requires_grad = ____
Show answer
for p in govde.parameters():
p.requires_grad = True
-
Match the stages.
- preparation
- Body learns on a side task
- stage 1
- Body frozen, head trained
- stage 2
- Everything, small lr
- parameter count
- Freezing check
Next lesson: Advanced section exam
