🎓 A1Learner
🔍

PyTorch The SGD Optimizer

SGD is an optimizer that moves each parameter a small step against its gradient. The name is short for stochastic gradient descent, and stochastic means random.

Creating SGD

SGD is created from the model's parameters and a learning rate. The stored rate can be read back from the param groups.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
print(optimizer.param_groups[0]["lr"])

The three-line rhythm

Every training round follows the same three lines in the same order.

optimizer.zero_grad()
loss.backward()
optimizer.step()

Ten steps of SGD

Repeating the rhythm in a loop brings the loss down fast.

import torch

torch.manual_seed(0)
x = torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[1.0], [3.0], [5.0], [7.0], [9.0]])
model = torch.nn.Linear(1, 1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

for step in range(10):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

loss = loss_fn(model(x), y)
print(f"Loss after 10 steps: {loss.item():.4f}")

Example

import torch

torch.manual_seed(0)
x = torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[1.0], [3.0], [5.0], [7.0], [9.0]])
model = torch.nn.Linear(1, 1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

for step in range(80):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

print(f"Final loss: {loss.item():.4f}")
print(f"Learned weight: {model.weight.item():.2f}")
print(f"Learned bias: {model.bias.item():.2f}")

Key points

  • Use torch.optim.SGD(model.parameters(), lr=...) to create SGD
  • Param groups: the optimizer's stored groups of settings, learning rate included
  • Use optimizer.zero_grad() to clear gradients before each backward()
  • Use loss.backward() to fill in fresh gradients
  • Use optimizer.step() to move the parameters once
  • SGD repeats one rhythm, and the loss falls a little on most rounds

In this example, SGD fits the same line rule, y = 2x + 1. The seed fixes the model's starting values. Each of the 80 rounds follows the same rhythm. The round clears gradients, computes the loss, runs backward, and steps. No update line is written by hand anywhere in the loop. The final loss prints as 0.0001. The learned weight and bias land close to the true 2 and 1.

Note

Forgetting zero_grad() lets gradients add up, so the next step uses stale gradients too

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Which order do these three calls follow in one training round?