Automatic Differentiation with PyTorch Autograd

Hello again! Welcome to our second lesson on PyTorch.

In our last lesson, we explored torch.Tensor, the fundamental building block of PyTorch. We saw how similar it is to a NumPy array but hinted at its "superpowers." Before that, in our NumPy module, you meticulously derived and implemented the backpropagation algorithm from scratch. You calculated gradients using the chain rule, updated weights manually, and truly understood the mechanics of how a neural network learns.

Today, we'll connect these two experiences and unveil one of PyTorch's most powerful features: Autograd, its automatic differentiation engine. This is the "magic" that frees us from manual gradient calculations, allowing us to build and train complex models with ease. After this lesson, you will understand exactly how PyTorch automates the work you did by hand.

Our goal is to learn how to use PyTorch's autograd to compute gradients automatically.

1. The Computational Graph: How PyTorch Tracks Operations

The core of autograd is the computational graph. Conceptually, this is a directed acyclic graph (DAG) where nodes represent operations (like addition or multiplication) and edges represent the flow of tensors. When you perform operations on tensors, PyTorch dynamically builds this graph in the background.

This process only happens for tensors that you've explicitly marked as needing gradients. These are typically the learnable parameters of your model (weights and biases). You tell PyTorch to track operations on a tensor by setting its requires_grad attribute to True.

Any tensor created as a result of an operation on a requires_grad tensor will also be part of this graph. These new tensors will have a grad_fn (gradient function) attribute, which references the operation that created them. The tensors you create directly (like your weights w and bias b) are called leaf nodes of the graph, and it's their gradients we are ultimately interested in.

Let's solidify this with a textual explanation from the official PyTorch documentation.

2. Computing Gradients with .backward()

Once PyTorch has built the computational graph during the forward pass (where you compute your output and loss), you can automatically compute all gradients using a single command: loss.backward().

When you call .backward() on a scalar tensor (like your final loss value), PyTorch does the following:

1. It traverses the computational graph backward from that scalar.
2. It uses the chain rule to compute the gradient of the scalar with respect to every tensor in the graph.
3. It stores the computed gradients in the .grad attribute of the leaf nodes (i.e., the tensors you created with requires_grad=True).

This process is the backward pass. It does exactly what you did manually in the NumPy module, but it's fully automated.

Let's walk through a simple example. Consider the function $y=w\cdot x$ and a squared error loss $L=(y-y_{true}{)}^2$. The gradient of the loss with respect to the weight w is $\frac{\partial L}{\partial w}=2(wx-y_{true})\cdot x$.

Let's see autograd compute this for us.

import torch




# Setup parameters and data
w = torch.tensor([3.0], requires_grad=True) # Our learnable parameter
x = torch.tensor([5.0])                      # Input data
y_true = torch.tensor([12.0])                # Target value




# 1. Forward pass
y = w * x                                    # Prediction: y = 3.0 * 5.0 = 15.0
loss = (y - y_true)**2                       # Loss: (15.0 - 12.0)^2 = 3.0^2 = 9.0




# 2. Backward pass
loss.backward()




# 3. Inspect the gradient
print(w.grad)



# Expected analytical gradient: 2 * (3.0*5.0 - 12.0) * 5.0 = 2 * 3.0 * 5.0 = 30.0
# PyTorch output: tensor([30.])

As you can see, PyTorch computed the exact gradient that you would have derived using calculus.

3. Essential Practices for Training

While .backward() is powerful, there are two crucial details you must manage in every training loop to ensure correctness.

a) Gradient Accumulation and Zeroing

By default, whenever you call .backward(), PyTorch accumulates (sums up) the new gradients into the .grad attributes. It does not overwrite them. This can be useful for advanced techniques, but for standard training, it leads to incorrect updates.

Therefore, before each backward pass in your training loop, you must manually zero out the gradients from the previous step.

b) Disabling Gradient Tracking

When you are not training—for example, when running your model for inference or evaluating its performance—you should disable gradient tracking. This is because the tracking mechanism adds computational overhead. Disabling it makes your code run faster. The most common way to do this is with the with torch.no_grad(): context manager.

The following video segments clearly illustrate these two essential practices.

Here's a summary of the common workflow in a training step:

1. Zero previous gradients: optimizer.zero_grad() or model.zero_grad().
2. Forward pass: prediction = model(inputs).
3. Compute loss: loss = loss_function(prediction, targets).
4. Backward pass: loss.backward().
5. Update weights: optimizer.step().

We will formalize this loop in the coming lessons.

4. Optional: Gradients of Non-Scalar Outputs

So far, we have only called .backward() on a scalar (the loss). What happens if you call it on a tensor with multiple elements? In this case, autograd computes something called a vector-Jacobian product, not the full gradient matrix (the Jacobian), for efficiency. This is a more advanced topic, but given your math background, it might be an interesting detail. It requires passing a vector of the same shape as the output tensor to the .backward() call.

Conclusion

In this lesson, you've bridged the gap between the manual backpropagation you implemented in NumPy and the automated, efficient world of PyTorch. autograd is a cornerstone of modern deep learning frameworks, and you now understand how it works.

Key Takeaways:

• PyTorch builds a dynamic computational graph to track all operations on tensors where requires_grad=True.
• Calling loss.backward() on a scalar loss automatically computes the gradients for all learnable parameters (leaf nodes) using the chain rule.
• The computed gradients are stored in the .grad attribute of each parameter tensor.
• Gradients accumulate by default. You must call .grad.zero_() (or optimizer.zero_grad()) before each new backward pass in a training loop.
• For inference or evaluation, you should disable gradient tracking using with torch.no_grad(): to improve performance.

Now that you can get gradients automatically, the next step is to structure our parameters and operations into clean, reusable components. In the next lesson, we will learn how to define a complete neural network layer and model using PyTorch's nn.Module and nn.Linear.