Challenge 3
- Machine Learning course
- Advanced
- Lesson 77/78
- 8 min
Build an image model with CNN.
Challenge 3: your own optimiser
The best way to prove what is inside PyTorch's optimisers is to write one yourself. A class deriving from torch.optim.Optimizer walks the parameter groups inside step:
class BenimSGD(torch.optim.Optimizer):
def __init__(self, params, lr):
super().__init__(params, {"lr": lr})
@torch.no_grad()
def step(self):
for grup in self.param_groups:
for p in grup["params"]:
if p.grad is not None:
p -= grup["lr"] * p.gradzero_grad, parameter groups and state_dict come ready from the base class.
Practice MySGD
Write the BenimSGD class from the card. Two model copies from the same start are given; train one with BenimSGD(lr=0.1) and the other with torch.optim.SGD(lr=0.1) for 20 steps. Print with torch.allclose that the final weights are the same.
Starter code
import copy
import torch
from torch import nn
torch.manual_seed(0)
X = torch.randn(64, 3)
y = X @ torch.tensor([[1.0], [-2.0], [0.5]])
a = nn.Linear(3, 1)
b = copy.deepcopy(a)
# buraya yazSample solution
import copy
import torch
from torch import nn
torch.manual_seed(0)
X = torch.randn(64, 3)
y = X @ torch.tensor([[1.0], [-2.0], [0.5]])
a = nn.Linear(3, 1)
b = copy.deepcopy(a)
class BenimSGD(torch.optim.Optimizer):
def __init__(self, params, lr):
super().__init__(params, {"lr": lr})
@torch.no_grad()
def step(self):
for grup in self.param_groups:
for p in grup["params"]:
if p.grad is not None:
p -= grup["lr"] * p.grad
for model, opt in [(a, BenimSGD(a.parameters(), lr=0.1)), (b, torch.optim.SGD(b.parameters(), lr=0.1))]:
for _ in range(20):
opt.zero_grad()
nn.functional.mse_loss(model(X), y).backward()
opt.step()
print(torch.allclose(a.weight, b.weight))Expected output
True
Practice A momentum version
Write the BenimMomentum class: keep a buffer self.state[p]["buf"] for every parameter; on the first step buf = g.clone(), then buf = m·buf + g; the update is p -= lr·buf. Take 2 steps for (w - 3)² from w = 0 with lr = 0.1, m = 0.9 and print w with round(..., 3) (expected: 1.62).
Starter code
import torch
w = torch.tensor(0.0, requires_grad=True)
# buraya yazSample solution
import torch
w = torch.tensor(0.0, requires_grad=True)
class BenimMomentum(torch.optim.Optimizer):
def __init__(self, params, lr, m):
super().__init__(params, {"lr": lr, "m": m})
@torch.no_grad()
def step(self):
for grup in self.param_groups:
for p in grup["params"]:
if p.grad is None:
continue
durum = self.state[p]
if "buf" not in durum:
durum["buf"] = p.grad.clone()
else:
durum["buf"] = grup["m"] * durum["buf"] + p.grad
p -= grup["lr"] * durum["buf"]
opt = BenimMomentum([w], lr=0.1, m=0.9)
for _ in range(2):
opt.zero_grad()
((w - 3) ** 2).backward()
opt.step()
print(round(w.item(), 3))Expected output
1.62
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 Why is <code>step</code> marked with <code>@torch.no_grad()</code>?
- It isn't required
- So the update isn't recorded in the computational graph Correct answer
- To make it faster
- To zero the gradients
Explanation: The parameters are changed in place; this isn't an operation to differentiate.
2 Pass the default settings to the parent class.
super().__init__(params, ____)super().__init__(params, {"lr": lr})Explanation: The second argument is the dict of default group settings.
3 Match the pieces.
| Term | Meaning |
|---|---|
param_groups |
Parameter groups with settings |
@torch.no_grad() |
Untracked update |
p.grad is not None |
Skip frozen parameters |
self.state[p] |
Per-parameter state |