🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Optimizer State

Optimizer state is the extra data an optimizer keeps beside the parameters. Adam's state holds two running averages for every parameter, plus a step count.

Empty state before a step

A new optimizer has seen no gradients yet. Its state for one weight is an empty dict.

import torch

torch.manual_seed(0)
model = torch.nn.Linear(4, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
print(optimizer.state[model.weight])

The keys after one step

One step creates three entries for the weight. The step entry counts the steps that parameter has taken.

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()
print(sorted(optimizer.state[model.weight].keys()))

The saved form

state_dict() bundles the state with the param groups. The group's learning rate rides along inside the same dict.

import torch

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

Example

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(4, 8),
    torch.nn.ReLU(),
    torch.nn.Linear(8, 2),
)
x = torch.randn(5, 4)
y = torch.tensor([0, 1, 0, 1, 1])
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

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

print(f"Loss: {loss.item():.4f}")
state = optimizer.state_dict()
print(f"Param groups: {len(state['param_groups'])}")
print(f"State entries: {len(state['state'])}")
for name, param in model.named_parameters():
    entry = optimizer.state[param]
    print(
        f"{name}: param {tuple(param.shape)}, "
        f"exp_avg {tuple(entry['exp_avg'].shape)}, "
        f"exp_avg_sq {tuple(entry['exp_avg_sq'].shape)}, "
        f"step {int(entry['step'].item())}"
    )

Key points

  • State: the per-parameter data an optimizer saves between steps
  • exp_avg: Adam's running average of the recent gradients
  • exp_avg_sq: Adam's running average of the squared gradients
  • Use optimizer.state[param] to read one parameter's state
  • Use optimizer.state_dict() to collect the full state in one dict
  • The state is Adam's memory, and it holds one entry for each parameter

In this example, a two-layer model takes one Adam step on five rows. The seed is set first, so the start is repeatable. The loss prints as 0.7177 before the step is applied. The state_dict holds one param group and four state entries. Each entry matches one named parameter in the printout. Every exp_avg shape equals its parameter shape. Every exp_avg_sq shape matches too. Every step count reads 1, because each parameter moved once.

Note

A fresh optimizer has empty state, and entries appear only after the first step()

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does Adam's state hold for each parameter?