🎓 A1Learner
🔍
0 of 79 lessons completed

PyTorch The Adam Optimizer

Adam is an optimizer that adjusts the step size for each parameter on its own. It reads each parameter's recent gradients to make that adjustment.

Creating Adam

Adam is created the same way as SGD, from parameters and a learning rate.

import torch

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

What Adam remembers

Adam keeps two running averages for each parameter. Their names show up in the optimizer's state after one step.

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.Adam(model.parameters(), lr=0.05)

loss = loss_fn(model(x), y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

state = optimizer.state[model.weight]
print(sorted(state.keys()))

Adam on uneven data

The second column here is ten times bigger than the first. Adam still brings the loss down in 20 steps.

import torch

torch.manual_seed(0)
x = torch.tensor([[0.0, 0.0], [1.0, 10.0], [2.0, 20.0],
                  [3.0, 30.0], [4.0, 40.0]])
y = torch.tensor([[1.0], [5.0], [9.0], [13.0], [17.0]])
model = torch.nn.Linear(2, 1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

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

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

Example

import torch

x = torch.tensor([[0.0, 0.0], [1.0, 10.0], [2.0, 20.0],
                  [3.0, 30.0], [4.0, 40.0]])
y = torch.tensor([[1.0], [5.0], [9.0], [13.0], [17.0]])
loss_fn = torch.nn.MSELoss()


def train(make_optimizer):
    torch.manual_seed(0)
    model = torch.nn.Linear(2, 1)
    optimizer = make_optimizer(model)
    losses = []
    for step in range(61):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        if step % 20 == 0:
            losses.append(f"{loss.item():.4f}")
        loss.backward()
        optimizer.step()
    return losses


sgd_losses = train(
    lambda model: torch.optim.SGD(model.parameters(), lr=0.001))
adam_losses = train(
    lambda model: torch.optim.Adam(model.parameters(), lr=0.05))

print(f"SGD losses: {', '.join(sgd_losses)}")
print(f"Adam losses: {', '.join(adam_losses)}")

Key points

  • Use torch.optim.Adam(model.parameters(), lr=...) to create Adam
  • Running average: an average that updates a little with each new value
  • Adam keeps two running averages of the gradients for every parameter
  • Use the same zero_grad, backward, and step rhythm with Adam as with SGD
  • Adam gives each parameter its own effective step size, instead of one size for all

In this example, both optimizers fit the same two-column data. The targets follow 2 times the first column, plus 0.2 times the second, plus 1. The second column runs from 0 to 40, and the first runs from 0 to 4. The train function seeds the model before each race, so both starts match. SGD needs a tiny rate here, or the big column makes it explode. Its loss stalls near 0.7678. Adam uses a bigger rate safely and keeps falling. Its loss reaches 0.0164 after 60 steps.

Note

Adam is a common first choice, and SGD can still win on small, well-scaled problems

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why can Adam use a bigger learning rate than SGD on data with uneven column sizes?