🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Gradient Clipping

Gradient clipping is a cap on the total size of the gradients in one step. clip_grad_norm_ shrinks all the gradients together when the total passes the cap.

One wild batch, no clipping

The model first trains to a loss of 0.0025 on clean data. One spike batch then wrecks it in a single unclipped 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.SGD(model.parameters(), lr=0.05)
for step in range(30):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
print(f"Before spike: {loss_fn(model(x), y).item():.4f}")
spike_x = torch.tensor([[100.0]])
spike_y = torch.tensor([[-100.0]])
optimizer.zero_grad()
spike_loss = loss_fn(model(spike_x), spike_y)
spike_loss.backward()
optimizer.step()
print(f"After spike: {loss_fn(model(x), y).item():.4f}")

The norm before the cap

clip_grad_norm_ reports the total norm it found. The spike batch pushes the norm to 59614.53, far past a cap of 1.0.

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(30):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
spike_x = torch.tensor([[100.0]])
spike_y = torch.tensor([[-100.0]])
optimizer.zero_grad()
spike_loss = loss_fn(model(spike_x), spike_y)
spike_loss.backward()
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
print(f"{total_norm.item():.2f}")

The same spike, clipped

Clipping holds the spike step to a small move. The clean-data loss stays low after the clipped 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.SGD(model.parameters(), lr=0.05)
for step in range(30):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
spike_x = torch.tensor([[100.0]])
spike_y = torch.tensor([[-100.0]])
optimizer.zero_grad()
spike_loss = loss_fn(model(spike_x), spike_y)
spike_loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
print(f"After clipped spike: {loss_fn(model(x), y).item():.4f}")

Example

import torch


def train(use_clip):
    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(30):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
    before = loss_fn(model(x), y).item()

    spike_x = torch.tensor([[100.0]])
    spike_y = torch.tensor([[-100.0]])
    optimizer.zero_grad()
    spike_loss = loss_fn(model(spike_x), spike_y)
    spike_loss.backward()
    total_norm = sum(
        param.grad.norm().item() ** 2 for param in model.parameters()
    ) ** 0.5
    if use_clip:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    after = loss_fn(model(x), y).item()
    return before, total_norm, after


before, norm, after_plain = train(False)
_, _, after_clipped = train(True)
print(f"Loss before spike: {before:.4f}")
print(f"Spike gradient norm: {norm:.2f}")
print(f"Loss after spike, no clipping: {after_plain:.4f}")
print(f"Loss after spike, clipped: {after_clipped:.4f}")

Key points

  • Gradient norm: the total size of all gradients together, as one number
  • Use torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) to clip
  • clip_grad_norm_ returns the total norm from before clipping
  • Clipping scales every gradient by one factor, so their directions stay unchanged
  • The cap limits how much damage one wild batch can do in a single step

In this example, one train function runs the spike test twice from the same seed. Each run first trains for 30 steps on the clean line data. The clean loss prints as 0.0025 before the spike. The spike batch holds one row with input 100 and target -100. Its gradient norm prints as 59614.53. Without clipping, one step sends the clean loss to 53659348.0000. With a cap of 1.0, the same step leaves the loss at 0.0185.

Note

Clipping runs after backward() and before step(), or it has nothing to cap

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does clip_grad_norm_ do when the total gradient norm passes max_norm?