🎓 A1Learner
🔍

PyTorch Parameter Groups

Parameter groups are separate bundles of parameters inside one optimizer. Each group carries its own settings, like its own learning rate.

One rate for everything

One plain optimizer call gives every parameter the same rate. A fresh head often needs bigger steps than a settled body.

optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

Two groups, two rates

A list of dicts replaces the plain parameter list. Each dict names its parameters and its own rate.

import torch

torch.manual_seed(0)
body = torch.nn.Linear(2, 4)
head = torch.nn.Linear(4, 1)
optimizer = torch.optim.SGD(
    [
        {"params": body.parameters(), "lr": 0.01},
        {"params": head.parameters(), "lr": 0.1},
    ]
)
print([group["lr"] for group in optimizer.param_groups])

What each group holds

Each group here holds two tensors. One tensor is the layer's weight, and one is its bias.

import torch

torch.manual_seed(0)
body = torch.nn.Linear(2, 4)
head = torch.nn.Linear(4, 1)
optimizer = torch.optim.SGD(
    [
        {"params": body.parameters(), "lr": 0.01},
        {"params": head.parameters(), "lr": 0.1},
    ]
)
print([len(group["params"]) for group in optimizer.param_groups])

Example

import torch

torch.manual_seed(0)


class TwoPartModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.body = torch.nn.Linear(2, 4)
        self.head = torch.nn.Linear(4, 1)

    def forward(self, x):
        return self.head(torch.relu(self.body(x)))


x = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0],
                  [2.0, 1.0], [1.0, 2.0], [2.0, 2.0], [3.0, 1.0]])
y = torch.tensor([[1.0], [3.0], [4.0], [6.0],
                  [9.0], [8.0], [11.0], [12.0]])
model = TwoPartModel()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(
    [
        {"params": model.body.parameters(), "lr": 0.01},
        {"params": model.head.parameters(), "lr": 0.1},
    ]
)

body_before = torch.cat([p.detach().clone().flatten() for p in model.body.parameters()])
head_before = torch.cat([p.detach().clone().flatten() for p in model.head.parameters()])

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

body_after = torch.cat([p.detach().flatten() for p in model.body.parameters()])
head_after = torch.cat([p.detach().flatten() for p in model.head.parameters()])
body_change = (body_after - body_before).abs().max().item()
head_change = (head_after - head_before).abs().max().item()

print(f"Group 0 lr: {optimizer.param_groups[0]['lr']}")
print(f"Group 1 lr: {optimizer.param_groups[1]['lr']}")
print(f"Final loss: {loss.item():.4f}")
print(f"Body max change: {body_change:.4f}")
print(f"Head max change: {head_change:.4f}")

Key points

  • Param group: a dict with a params list and that group's settings
  • Use a list of dicts to create an optimizer with several groups
  • Use optimizer.param_groups to read the groups back in order
  • Body: the main part of a model, and head: the final layer that gives the answer
  • Groups let one optimizer train different parts at different speeds

In this example, a small model has a body layer and a head layer. The data follows the rule 3 times the first input, plus 2 times the second, plus 1. One SGD optimizer holds two groups with rates 0.01 and 0.1. The model trains for 50 full-batch steps. The final loss prints as 0.0196. The head's biggest change is 3.1800, and the body's is 0.9576. The head moved more in this run, and its rate was ten times bigger.

Note

Groups keep their creation order, and param_groups[0] is the first group built

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

How does one optimizer give a model's head a bigger learning rate than its body?