Computation Graphs and Backpropagation

Manual gradient code is usually short. A typical example builds a tiny model with one input, one weight, and one bias, passes the input through a sigmoid and a binary cross-entropy loss, and then asks PyTorch for the gradient of that loss with respect to one of the parameters. The whole script may be only a dozen lines, but a wrong result almost always traces back to one specific line. This article teaches you to follow the script line by line, watching how each variable gets its value, so that you can notice the moment the program's state stops matching what you expected.

Before any gradient is requested, the script has already built a chain of operations. The input x1 is multiplied by the weight w1, the bias b is added to produce z, the sigmoid of z gives a, and the binary cross-entropy of a against the target y gives the loss. In symbols, the chain is $z = x_1 w_1 + b$ followed by $a = \sigma(z)$, with the loss computed from $a$ and $y$. Each of these steps is recorded by PyTorch as it runs, and that recorded chain is what people mean by the computation graph.

The diagram shows the direction of data during the forward pass, from the inputs toward the loss. The parameters w1 and b are marked as requiring a gradient because the script creates them with requires_grad set to True. Tensors without that flag, such as the input x1 and the target y, are not the variables whose gradients this script asks for. Notice that the loss depends on every node on the path from the parameters to the loss, which is why the gradient of the loss with respect to a parameter can be assembled by walking that path backward.

The gradient of a loss with respect to a parameter answers the question of how the loss changes when that parameter changes. In this script, PyTorch is asked to compute that answer with the grad function, which is imported from torch.autograd. The function takes the loss and the variable you want to differentiate with respect to. The call grad(loss, b, retain_graph=True) asks how the loss changes as the bias b changes, and the result is stored in a variable named grad_L_b. The same pattern works for any parameter created with requires_grad set to True, such as the weight w1.

By default, PyTorch destroys the computation graph after calculating the gradients, in order to free memory. This is a sensible default. A graph can hold many intermediate values, and once the gradient has been computed, a standard training step usually does not need them again. The consequence is that the graph is no longer available for another gradient request unless you explicitly ask PyTorch to keep it.

The lower path of the diagram is the default. After the first request, the graph is freed, and a second request against the same loss may fail. The upper path is what happens when the first call keeps the graph. Because the call can keep the graph available, a later call can still reach the values it needs.

The script changes the default with retain_graph=True. This keyword argument tells PyTorch to keep the graph in memory after the call finishes. The source gives the reason directly: the script will reuse the same computation graph shortly, so it must stay available. Both the gradient with respect to the weight w1 and the gradient with respect to the bias b come from the same loss, and that shared loss is why the graph must survive the first request.

When a manual gradient script gives a result you did not expect, work through it in order rather than guessing. The source recommends three checkpoints, in this sequence: the setup, each grad call, and then the printed output. The flow below shows how to move through them.