Automatic Differentiation

One of the main reasons for using PyTorch in deep learning is its ability to automatically compute the gradient of a function. More specifically, PyTorch has a built-in differentiation engine, called torch.autograd, which tracks the operations performed on tensors. The history of computation is recorded in a data structure called a computational graph, which is then used to compute derivatives using the chain rule. In this notebook, we will explore the basics of automatic differentiation in PyTorch.

import torch

Computational graph

In order to get familiar with the concept of a computation graph, we will create one for the following function.

\[f(x) = \frac{1}{N}\sum_{i=1}^N \left[(x_i + 2)^2 + 3\right]\]

The vector \(x\) is the input of the function, so we will compute the gradient of \(f(x)\) with respect to \(x\). The gradient will be a tensor of the same shape as \(x\).

Note

A gradient is defined for a scalar-valued function. For a vector-valued function, the corresponding derivative is the Jacobian matrix, which is never computed explicitly in PyTorch.

Tensor with gradient

The first thing we have to do is to create the tensor that represents the vector \(x\). We will set requires_grad=True to track the operations on this tensor. This will allow us to compute the gradient of \(f(x)\) with respect to \(x\). Note that by default, tensors are created without tracking operations on them.

x = torch.tensor([0., 1, 2], requires_grad=True)

Now let’s build the computation graph step by step. You can combine multiple operations in a single line, but we will separate them here to better understand how each operation is added to the computation graph.

a = x + 2
b = a ** 2
c = b + 3
y = c.mean()  # f(x) evaluated at the values stored in x

print(y)

tensor(12.6667, grad_fn=<MeanBackward0>)

The statements above create a computation graph that looks similar to the figure below. The nodes represent the intermediate tensors, and the edges represent the operations that transform those tensors. The arrows point from the outputs to the inputs of each operation, because this is the direction in which the gradient will be computed.

Backpropagation

We can perform backpropagation on the computation graph by calling backward() on the last output.

y.backward()

Reverse-mode autograd applies the chain rule from the scalar output back to every leaf tensor with requires_grad=True. For this example,

\[\frac{\partial f}{\partial x_i}=\frac{2(x_i+2)}{N}.\]

Note

A leaf tensor is a tensor created by the user and not the result of an operation.

Leaf tensors that require gradients have their .grad attribute populated. In our example, the leaf tensor x will have its .grad attribute filled with the gradient of \(f(x)\) with respect to the variable \(x\), evaluated at the actual value of x that we used to build the computation graph.

Intermediate non-leaf tensors such as a, b, and c participate in backpropagation but do not retain .grad by default; call .retain_grad() on an intermediate tensor before backward() only when that value is needed for debugging.

x.grad  # df/dx evaluated at the values stored in x

tensor([1.3333, 2.0000, 2.6667])

expected_grad = 2 * (x.detach() + 2) / x.numel()

assert torch.allclose(x.grad, expected_grad)

Note that if we recompute the tensor y and call y.backward() a second time, the gradient in x.grad will change. This happens because PyTorch accumulates the gradients in the .grad attribute, that is, it sums the new gradient with the old one.

a = x + 2
b = a ** 2
c = b + 3
y = c.mean()

y.backward()

x.grad

tensor([2.6667, 4.0000, 5.3333])

If we need to compute the gradient in a loop (e.g., to implement gradient descent), we have to zero out the gradient in each iteration before recomputing the output scalar tensor on which backward() is called.

Note

PyTorch’s reverse-mode autograd does not construct the full Jacobian matrices of intermediate tensors when backward() is called. Instead, it computes a vector–Jacobian product (VJP), where the “vector” is the upstream gradient. For a scalar output, the upstream vector reduces to the scalar \(1\), so this product is simply the usual gradient. For a non-scalar output, the user must supply an upstream tensor with the same shape as the output.

Disabling gradient tracking

By default, all tensors with requires_grad=True are tracking their computational history and support gradient computation. However, there are some cases when we do not need to do that, for example, when we have trained a neural network and just want to apply it to some input data.

We can stop tracking computations by surrounding our computation code with a torch.no_grad() block.

x = torch.randn(5, requires_grad=True)

# Gradient is tracked
a = x * 2

# Tracking is stopped in the block
with torch.no_grad():
    z = x * 2

print("Does 'x' have a gradient?", x.requires_grad)
print("Does 'a' have a gradient?", a.requires_grad)
print("Does 'z' have a gradient?", z.requires_grad)

Does 'x' have a gradient? True
Does 'a' have a gradient? True
Does 'z' have a gradient? False

The function torch.no_grad() disables recording inside a block. The method detach() returns a tensor disconnected from the graph but sharing storage with its source, so in-place mutation can still affect both. Use detach().clone() for independent storage. For evaluation, torch.inference_mode() goes further than no_grad() by disabling additional autograd bookkeeping; tensors created there are more restrictive if later reused in gradient-tracked computations.

a = x * 2
z = a.detach()

print("Does 'a' have a gradient?", a.requires_grad)
print("Does 'z' have a gradient?", z.requires_grad)

Does 'a' have a gradient? True
Does 'z' have a gradient? False

Optimization

After understanding the fundamentals of autograd in PyTorchs, the next logical step is to apply this knowledge to practical optimization problems. One common task in machine learning and numerical analysis is finding the minimum of a function, a process that can be achieved using an optimization algorithm called gradient descent. PyTorch provides a variety of optimization algorithms that make this process efficient and straightforward. These algorithms are implemented in the torch.optim module.

Loss function

In order to apply gradient descent, we need to define a function that we want to minimize. This function should take a PyTorch tensor as input and return a scalar output. For example, let’s consider the following function of one variable.

\[ f(x) = x^4 + x^2 + 10x \]

We can define this function in PyTorch as follows.

def function(x):
    return x ** 4 + x ** 2 + 10 * x

The goal is to find the value of \(x\) that minimizes this expression.

Initialization

The optimization process requires an initial guess. Hence, we need to create a tensor representing the initial value of \(x\) and ensure it requires a gradient.

x = torch.tensor(0.0, requires_grad=True)

Choosing an optimizer

PyTorch provides several optimization algorithms in torch.optim. We pass the leaf tensors to optimize and a learning rate. Initialization and learning rate affect whether optimization converges and which minimum is reached; in neural networks, parameters are normally initialized by their layer modules. Here we use SGD with a learning rate of 0.05.

optimizer = torch.optim.SGD([x], lr=0.05)

All optimizers implement a step() method that updates the provided tensors based on the computed gradients, and a zero_grad() method that resets those gradients to zero.

Running the optimization loop

The next step is to perform the optimization loop. Within each iteration, the following steps are performed.

• Clear gradients accumulated during the previous iteration.

• Compute the value of the function by calling it with the current x value.

• Perform backpropagation by invoking the backward() method on the computed function value. This populates the gradient in x.grad.

• Update x by calling optimizer.step(), which applies a single step of the optimization algorithm.

The loop can be run for a fixed number of steps or until the change in the function value is sufficiently small.

for i in range(10):
    optimizer.zero_grad()
    loss = function(x)
    loss.backward()
    optimizer.step()

    with torch.no_grad():
        print(f"Iteration {i+1:2d}: x = {x.item():.6f}, f(x) = {function(x).item():.7f}")

Iteration  1: x = -0.500000, f(x) = -4.6875000
Iteration  2: x = -0.925000, f(x) = -7.6622810
Iteration  3: x = -1.174209, f(x) = -8.4623260
Iteration  4: x = -1.232996, f(x) = -8.4984322
Iteration  5: x = -1.234797, f(x) = -8.4984636
Iteration  6: x = -1.234772, f(x) = -8.4984646
Iteration  7: x = -1.234773, f(x) = -8.4984636
Iteration  8: x = -1.234773, f(x) = -8.4984636
Iteration  9: x = -1.234773, f(x) = -8.4984636
Iteration 10: x = -1.234773, f(x) = -8.4984636

Visualizing the result

Finally, we can visualize the function and the minimum value found by the optimization algorithm.

import numpy as np
import matplotlib.pyplot as plt

t = np.linspace(-2.5, 1.5, 100)
f = function(torch.tensor(t)).numpy()
plt.plot(t, f)
plt.plot(x.item(), function(x).item(), 'ro')
plt.show()

Conclusion

In this notebook, we have explored the basics of automatic differentiation in PyTorch. We have seen how to create a computation graph, compute gradients, and perform optimization using the built-in optimization algorithms. This knowledge is essential for understanding how PyTorch works.