PyTorch SGD vs Adam in Practice

SGD and Adam are two optimizers that update parameters from gradients. SGD uses the current gradient, and Adam rescales each step from recent gradients.

SGD with momentum

Momentum keeps a share of the past gradients in each new step. The step then points along the recent path, not just the newest gradient.

optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

Plain SGD after 20 steps

The blob data has three round clouds of points. Plain SGD brings the loss to 0.1775 in 20 steps.

import torch

torch.manual_seed(0)
centers = torch.tensor([[-2.0, -2.0], [2.0, -2.0], [0.0, 2.0]])
parts = []
labels = []
for label, center in enumerate(centers):
    parts.append(center + torch.randn(30, 2) * 0.6)
    labels.extend([label] * 30)
x = torch.cat(parts)
y = torch.tensor(labels)
loss_fn = torch.nn.CrossEntropyLoss()

torch.manual_seed(1)
model = torch.nn.Linear(2, 3)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for step in range(20):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
with torch.no_grad():
    print(f"{loss_fn(model(x), y).item():.4f}")

Adam and momentum after 20 steps

Adam reaches 0.1035 in the same 20 steps. Momentum SGD reaches 0.0089, the lowest of the three at this point.

import torch

torch.manual_seed(0)
centers = torch.tensor([[-2.0, -2.0], [2.0, -2.0], [0.0, 2.0]])
parts = []
labels = []
for label, center in enumerate(centers):
    parts.append(center + torch.randn(30, 2) * 0.6)
    labels.extend([label] * 30)
x = torch.cat(parts)
y = torch.tensor(labels)
loss_fn = torch.nn.CrossEntropyLoss()


def train20(make_optimizer):
    torch.manual_seed(1)
    model = torch.nn.Linear(2, 3)
    optimizer = make_optimizer(model)
    for step in range(20):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        return loss_fn(model(x), y).item()


adam_loss = train20(lambda m: torch.optim.Adam(m.parameters(), lr=0.05))
momentum_loss = train20(
    lambda m: torch.optim.SGD(m.parameters(), lr=0.1, momentum=0.9))
print(f"Adam: {adam_loss:.4f}")
print(f"Momentum SGD: {momentum_loss:.4f}")

Example

import torch

torch.manual_seed(0)
centers = torch.tensor([[-2.0, -2.0], [2.0, -2.0], [0.0, 2.0]])
parts = []
labels = []
for label, center in enumerate(centers):
    parts.append(center + torch.randn(30, 2) * 0.6)
    labels.extend([label] * 30)
x = torch.cat(parts)
y = torch.tensor(labels)
loss_fn = torch.nn.CrossEntropyLoss()


def train(make_optimizer):
    torch.manual_seed(1)
    model = torch.nn.Linear(2, 3)
    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()
    with torch.no_grad():
        final_loss = loss_fn(model(x), y).item()
        accuracy = (model(x).argmax(dim=1) == y).float().mean().item()
    losses.append(f"{final_loss:.4f}")
    return losses, accuracy


runs = [
    ("SGD", lambda m: torch.optim.SGD(m.parameters(), lr=0.1)),
    ("SGD + momentum", lambda m: torch.optim.SGD(m.parameters(), lr=0.1, momentum=0.9)),
    ("Adam", lambda m: torch.optim.Adam(m.parameters(), lr=0.05)),
]
for name, make_optimizer in runs:
    losses, accuracy = train(make_optimizer)
    print(f"{name}: losses {', '.join(losses)}, accuracy {accuracy:.2f}")

Key points

Momentum: a running average of past gradients added to the current step
Use momentum=0.9 to turn momentum on in SGD
Use the same seed before each run to give every optimizer the same start
On this task, momentum SGD reached the lowest final loss
All three runs reached accuracy 1.00 on this task
The three runs differ only in the optimizer, and the traces show what that one choice changes

In this example, three optimizers train the same classifier on the same blob data. The train function reseeds the model before each run, so every start matches. Each trace records the loss every 20 steps and once more at the end. Plain SGD ends at 0.0658, the highest final loss. Adam ends at 0.0217, lower than plain SGD here. Momentum SGD ends at 0.0029, the lowest in this race. All three runs still reach accuracy 1.00 on this easy task.

Note

A winner on one task is not a winner on every task, and new problems deserve a short race

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

torch.manual_seed(0)
centers = torch.tensor([[-2.0, -2.0], [2.0, -2.0], [0.0, 2.0]])
parts = []
labels = []
for label, center in enumerate(centers):
    parts.append(center + torch.randn(30, 2) * 0.6)
    labels.extend([label] * 30)
x = torch.cat(parts)
y = torch.tensor(labels)
loss_fn = torch.nn.CrossEntropyLoss()


def train(make_optimizer):
    torch.manual_seed(1)
    model = torch.nn.Linear(2, 3)
    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()
    with torch.no_grad():
        final_loss = loss_fn(model(x), y).item()
        accuracy = (model(x).argmax(dim=1) == y).float().mean().item()
    losses.append(f"{final_loss:.4f}")
    return losses, accuracy


runs = [
    ("SGD", lambda m: torch.optim.SGD(m.parameters(), lr=0.1)),
    ("SGD + momentum", lambda m: torch.optim.SGD(m.parameters(), lr=0.1, momentum=0.9)),
    ("Adam", lambda m: torch.optim.Adam(m.parameters(), lr=0.05)),
]
for name, make_optimizer in runs:
    losses, accuracy = train(make_optimizer)
    print(f"{name}: losses {', '.join(losses)}, accuracy {accuracy:.2f}")

Exercise

What does momentum add to plain SGD?

A running average of past gradients in each step A second loss function A bigger batch of data A copy of the best weights