Automatic gradients in PyTorch

PyTorch records each step on the way forward

Take a tiny model with one weight w. It predicts p = w × x for an input x = 2, and the right answer is t = 7. The error is e = p − t. The loss, one number that says how wrong the model is, is L = e². With w = 3, the steps give p = 6, e = −1 and L = 1.

When w is a tensor with requires_grad=True, PyTorch keeps a record of these steps as it runs them. The record is the computation graph: each result remembers which operation made it and from which inputs. The figure draws it.

.backward() multiplies local slopes from the end to the start

Each step in the graph has a simple local slope: how much its output changes per unit change of its input. PyTorch knows the local slope of every built-in operation. L.backward() starts at L, where the slope of L with respect to itself is 1. It then moves back one step at a time, multiplying by each local slope. That product is the chain rule. The result for each marked tensor is stored in its .grad.

one_weight.py

import torch

x, t = 2.0, 7.0
w = torch.tensor(3.0, requires_grad=True)
p = w * x          # prediction: 6
e = p - t          # error: -1
L = e ** 2         # loss: 1
L.backward()       # work backward from L
print(L.item(), w.grad.item())

1.0 -4.0

.item() turns a one-number tensor into a plain Python number.

One backward pass gives the gradient for every weight

Real models have millions of weights. A single call to .backward() fills .grad for all of them, and it costs about as much as a few forward passes. The list of these derivatives is the gradient, and it is what gradient descent uses to update the weights. Backpropagation shows how this backward pass runs through the layers of a network.

Here the model has three weights and one loss. Before you read the output, use the worked example to predict w.grad: the error is −2, so each gradient is 2 × (−2) × x.

three_weights.py

import torch

w = torch.tensor([1.0, -2.0, 0.5], requires_grad=True)
x = torch.tensor([3.0, 1.0, 4.0])
L = ((w * x).sum() - 5.0) ** 2   # one number from three weights
L.backward()
print(L.item())
print(w.grad)

4.0
tensor([-12.,  -4., -16.])

Two habits: clear old gradients, and switch recording off when you do not need it

.backward() adds to .grad; it does not replace it. Training code therefore clears the gradients before each step, by setting them to None or by calling the optimizer's zero_grad(). Recording also costs memory, because the graph keeps every intermediate result. Code that only uses the model, such as testing or generating, runs inside torch.no_grad(), which turns recording off.

habits.py

import torch

w = torch.tensor(3.0, requires_grad=True)
for step in range(2):
    L = (w * 2.0 - 7.0) ** 2
    L.backward()
    print("after backward", step + 1, "w.grad =", w.grad.item())
w.grad = None                  # clear it before the next step

with torch.no_grad():          # no recording inside this block
    L = (w * 2.0 - 7.0) ** 2
print(L.requires_grad)

after backward 1 w.grad = -4.0
after backward 2 w.grad = -8.0
False