Challenge 4
- Machine Learning course
- Advanced
- Lesson 78/78
- 8 min
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 togetherChecking the trainable parameter count at every stage verifies that the freezing really took effect.
Practice 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.
Starter code
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)
# buraya yazSample 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)Expected output
66 True
Practice 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.
Starter code
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
# buraya yazSample 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)Expected output
1218 True
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 Which parameters should be given to the optimiser in stage 1?
- Only the new head's parameters Correct answer
- All parameters
- Only the body
- None
Explanation: Since the body is frozen it won't be updated; only the head is given to the optimiser.
2 Make the body trainable again.
for p in govde.parameters():
p.requires_grad = ____for p in govde.parameters():
p.requires_grad = TrueExplanation: In stage 2 the body is opened.
3 Match the stages.
| Term | Meaning |
|---|---|
preparation |
Body learns on a side task |
stage 1 |
Body frozen, head trained |
stage 2 |
Everything, small lr |
parameter count |
Freezing check |