🎯 What You'll Learn
- Install PyTorch, verify CUDA availability, and understand the device model (CPU vs GPU)
- Create tensors in every important way: from lists, zeros/ones/eye, random, arange, linspace
- Perform tensor operations: element-wise math, matrix multiplication, reshaping, aggregations
- Understand the computation graph: what autograd tracks and why
- Compute your first gradient with
.backward()and verify it manually - Understand gradient accumulation and the critical
.zero_grad()pattern - Use
torch.no_grad()for efficient inference - Distinguish leaf tensors from non-leaf tensors and understand where
.gradlives
NumPy arrays are brilliant for numerical computation, but they have a critical limitation: they don't know how to compute gradients. PyTorch tensors are like NumPy arrays that went to graduate school — they can not only store data and do maths, but they also keep a diary of every operation they've been through, ready to compute gradients automatically via the chain rule. This "diary" is the autograd system, and it is the single most important innovation that made modern deep learning tractable. Backpropagation through millions of parameters in a single backward pass — that's what autograd does for you.
1 Installing PyTorch and First Steps
PyTorch is not included with a standard Python install. Installation depends on your platform and whether you have a CUDA-capable GPU.
# ── Installation ──
# CPU-only (works on any machine, fine for learning):
# pip install torch torchvision
# With CUDA GPU support (check pytorch.org/get-started for your CUDA version):
# pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
# On Google Colab: PyTorch is pre-installed with GPU support.
# ── First steps ──
import torch
import numpy as np
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
# Standard device setup — write this at the top of every DL project
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"\nUsing device: {device}")PyTorch Tensor vs NumPy Array: Side-by-Side
The surface-level syntax is very similar to NumPy, but there are three critical differences: tensors can live on GPU, tensors can track gradients, and tensor operations are integrated with PyTorch's neural network ecosystem.
import numpy as np
import torch
# ── Same operation in NumPy and PyTorch ──
# NumPy
a_np = np.array([1.0, 2.0, 3.0])
b_np = np.array([4.0, 5.0, 6.0])
c_np = a_np * b_np + 2.0
print(f"NumPy: {c_np} dtype: {c_np.dtype}") # [6. 12. 20.] float64
# PyTorch
a_pt = torch.tensor([1.0, 2.0, 3.0])
b_pt = torch.tensor([4.0, 5.0, 6.0])
c_pt = a_pt * b_pt + 2.0
print(f"PyTorch: {c_pt} dtype: {c_pt.dtype}") # [6., 12., 20.] float32
# Key differences
print(f"\nNumPy array type: {type(a_np)}") # numpy.ndarray
print(f"PyTorch tensor type: {type(a_pt)}") # torch.Tensor
print(f"NumPy default dtype: {a_np.dtype}") # float64
print(f"PyTorch default dtype: {a_pt.dtype}") # float32 ← ML standard
# The crucial difference: gradient tracking
x = torch.tensor(3.0, requires_grad=True) # "track me!"
y = x ** 2
print(f"\ny = x² at x=3: {y}") # tensor(9.)
y.backward() # compute dy/dx
print(f"dy/dx = 2x = {x.grad}") # tensor(6.) ✓
# NumPy can't do this
x_np = np.array(3.0)
# x_np.backward() ← AttributeError: no such methodNumPy defaults to float64 (64-bit, 8 bytes per number) for precision in scientific computing. PyTorch defaults to float32 (32-bit, 4 bytes per number) to halve GPU memory usage — crucial for large models. When you convert data between libraries, always check dtypes: torch.from_numpy(np_array) preserves the dtype, so a float64 NumPy array becomes a float64 PyTorch tensor. Neural networks almost always expect float32. Use .float() or .to(torch.float32) to convert.
2 Creating Tensors: Every Way You Need to Know
There are many functions for creating tensors. Knowing all of them saves time and prevents shape bugs. Here is the complete practical reference.
import torch
# ── From Python data ──
t1 = torch.tensor([1, 2, 3]) # 1D int tensor
t2 = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # 2D float tensor
t3 = torch.tensor([True, False, True]) # bool tensor
print(f"From list: {t2.shape}, dtype: {t2.dtype}") # (2,2), float32
# ── Zeros, ones, constants ──
zeros = torch.zeros(3, 4) # 3×4 tensor of 0.0
ones = torch.ones(2, 3) # 2×3 tensor of 1.0
full = torch.full((2, 3), 7.5) # 2×3 tensor of 7.5
eye = torch.eye(4) # 4×4 identity matrix
empty = torch.empty(2, 3) # uninitialized memory (fast, but random garbage)
print(f"zeros: {zeros.shape}, ones: {ones.shape}, eye: {eye.shape}")
# ── Random tensors ──
rand_uniform = torch.rand(3, 4) # uniform [0, 1)
rand_normal = torch.randn(100, 5) # standard normal N(0,1)
rand_int = torch.randint(0, 10, (4, 4)) # integers in [0, 10)
perm = torch.randperm(10) # random permutation of 0..9
torch.manual_seed(42) # reproducibility
x = torch.randn(3)
print(f"randn(3) with seed 42: {x.round(decimals=4)}")
# ── Sequences ──
arange = torch.arange(0, 10, 2) # [0, 2, 4, 6, 8]
linspace = torch.linspace(0, 1, 5) # [0.0, 0.25, 0.5, 0.75, 1.0]
print(f"arange: {arange}")
print(f"linspace: {linspace}")
# ── Dtype control ──
f32 = torch.tensor([1, 2, 3], dtype=torch.float32)
i64 = torch.tensor([1, 2, 3], dtype=torch.int64)
b = torch.tensor([True, False], dtype=torch.bool)
print(f"f32 dtype: {f32.dtype}, i64 dtype: {i64.dtype}")
# ── Inspecting properties ──
x = torch.randn(32, 3, 224, 224) # batch of 32 RGB images, ImageNet size
print(f"\nshape: {x.shape}") # torch.Size([32, 3, 224, 224])
print(f"dtype: {x.dtype}") # torch.float32
print(f"device: {x.device}") # cpu (or cuda:0)
print(f"ndim: {x.ndim}") # 4
print(f"numel: {x.numel():,}") # total elements: 32*3*224*224 = 4,816,896torch.zeros_like(t) creates a zero tensor with the same shape, dtype, and device as tensor t. This is extremely useful in practice: when computing gradients manually, initializing buffers for accumulation, or creating masks. torch.rand_like(t) and torch.randn_like(t) also exist. These convenience functions prevent shape mismatches — you can't accidentally create a gradient tensor with the wrong shape.
3 Tensor Operations: Math Without Loops
One of PyTorch's most important features is that almost all operations work on entire tensors without Python loops — making them fast (they call optimized C++/CUDA code underneath).
import torch
A = torch.tensor([[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0]]) # shape (2, 3)
B = torch.tensor([[7.0, 8.0, 9.0],
[1.0, 2.0, 3.0]]) # shape (2, 3)
# ── Element-wise operations ──
print(A + B) # [[8, 10, 12], [5, 7, 9]]
print(A * B) # [[7, 16, 27], [4, 10, 18]] (element-wise, NOT matmul)
print(A ** 2) # [[1, 4, 9], [16, 25, 36]]
print(torch.sqrt(A)) # element-wise square root
# ── Matrix multiplication ──
# For matmul, shapes must be compatible: (m,k) @ (k,n) = (m,n)
A_sq = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # (2,2)
B_sq = torch.tensor([[5.0, 6.0], [7.0, 8.0]]) # (2,2)
C1 = torch.matmul(A_sq, B_sq) # official function
C2 = A_sq @ B_sq # @ operator (cleaner, identical result)
print(f"\nMatrix multiply:\n{C1}")
# [[19, 22],
# [43, 50]]
# Batched matmul (crucial for DL — whole batch at once)
batch_A = torch.randn(32, 64, 128) # 32 matrices, each 64×128
batch_W = torch.randn(32, 128, 256) # 32 matrices, each 128×256
result = torch.bmm(batch_A, batch_W) # batched matmul
print(f"Batched matmul: {result.shape}") # (32, 64, 256)
# ── Transpose ──
x = torch.randn(3, 4)
print(f"\nx shape: {x.shape}") # (3, 4)
print(f"x.T shape: {x.T.shape}") # (4, 3)
print(f"transpose(0,1): {x.transpose(0, 1).shape}") # (4, 3) — same thing
# ── Aggregations ──
x = torch.tensor([[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0]])
print(f"\nSum all: {x.sum()}") # 21.0
print(f"Sum rows: {x.sum(dim=1)}") # [6., 15.]
print(f"Sum cols: {x.sum(dim=0)}") # [5., 7., 9.]
print(f"Mean: {x.mean():.4f}") # 3.5
print(f"Max: {x.max()}, Min: {x.min()}")
print(f"Argmax (flat): {x.argmax()}") # 5 (index of 6.0)
print(f"Argmax (rows): {x.argmax(dim=1)}") # [2, 2] (max in each row)
# ── Reshaping ──
flat = torch.arange(12.0) # [0, 1, 2, ..., 11]
matrix = flat.reshape(3, 4) # shape (3, 4)
matrix_inferred = flat.reshape(-1, 4) # same: -1 means "infer this dim"
col_vec = flat.reshape(-1, 1) # shape (12, 1)
row_vec = flat.reshape(1, -1) # shape (1, 12)
# view vs reshape: view requires contiguous memory; reshape handles both
x = torch.randn(2, 3, 4)
y = x.view(6, 4) # shape (6, 4) — shares memory with x
z = x.reshape(6, 4) # shape (6, 4) — may copy if needed
# unsqueeze / squeeze: add or remove dimensions of size 1
x = torch.randn(3) # shape (3,)
x_col = x.unsqueeze(1) # shape (3, 1) — column vector
x_row = x.unsqueeze(0) # shape (1, 3) — row vector
x_back = x_col.squeeze() # shape (3,) — removes all size-1 dims
print(f"\nOriginal: {x.shape}, unsqueeze(1): {x_col.shape}, squeeze: {x_back.shape}")Moving Between Devices and Libraries
import torch
import numpy as np
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Create on CPU, move to GPU
t_cpu = torch.randn(100, 100)
t_gpu = t_cpu.to(device) # move to GPU (no-op if already there)
t_gpu2 = t_cpu.cuda() # shortcut for GPU (only if CUDA available)
t_back = t_gpu.cpu() # move back to CPU
# NumPy bridge (CPU only — GPU tensors must be .cpu() first)
np_arr = np.array([1.0, 2.0, 3.0])
pt_tensor = torch.from_numpy(np_arr) # shares memory! (no copy)
back_to_np = pt_tensor.numpy() # back to NumPy (shares memory)
# If tensor tracks gradients, must detach first
x = torch.tensor([1.0, 2.0], requires_grad=True)
x_np = x.detach().numpy() # detach, then convert
print(f"Tensor to numpy: {x_np}")
# Common dtype conversions
t = torch.tensor([1, 2, 3]) # int64 by default
t_float = t.float() # → float32
t_double = t.double() # → float64
t_half = t_float.half() # → float16 (for mixed precision)
print(f"dtypes: {t.dtype}, {t_float.dtype}, {t_double.dtype}")4 The Computation Graph: What Autograd Tracks
This is where PyTorch becomes different from NumPy. When you perform operations on tensors that have requires_grad=True, PyTorch silently builds a computation graph in the background — a record of every operation you performed, ready to compute derivatives using the chain rule.
The Diary Analogy in Detail
Imagine PyTorch keeps a diary while you compute. Every entry says: "I used operation X to produce tensor Y from tensor Z." When you later call y.backward(), PyTorch reads the diary backwards, applying the chain rule at each entry to compute the gradient of the final output with respect to each parameter that contributed.
The diary is the computation graph — technically a Directed Acyclic Graph (DAG) where each node is an operation and each edge is a tensor. The graph is built dynamically as your code runs (this is what "define-by-run" means).
import torch
# ── See the computation graph in action ──
x = torch.tensor(2.0, requires_grad=True) # "track gradients for x"
y = torch.tensor(3.0, requires_grad=True) # "track gradients for y"
# Step 1: Multiply
z = x * y # z = xy
print(f"z = x*y = {z.item()}") # 6.0
print(f"z.requires_grad: {z.requires_grad}") # True (inherited)
print(f"z.grad_fn: {z.grad_fn}") # MulBackward0 — records the * op
# Step 2: Add a constant
w = z + 4 # w = xy + 4
print(f"w.grad_fn: {w.grad_fn}") # AddBackward0
# Step 3: Square
out = w ** 2 # out = (xy + 4)²
print(f"out.grad_fn: {out.grad_fn}") # PowBackward0
# The chain is: out → PowBackward(w) → AddBackward(z, 4) → MulBackward(x, y)
# This IS the computation graph.
# ── requires_grad propagation rules ──
a = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(2.0, requires_grad=False)
c = a + b # c inherits requires_grad=True (because a does)
d = b + b # d is False (neither parent has requires_grad=True)
print(f"\na.requires_grad: {a.requires_grad}") # True
print(f"b.requires_grad: {b.requires_grad}") # False
print(f"c.requires_grad: {c.requires_grad}") # True (a contributes)
print(f"d.requires_grad: {d.requires_grad}") # False
# ── is_leaf: was this tensor created directly or by an operation? ──
x = torch.tensor(5.0, requires_grad=True)
y = x * 2 # non-leaf: produced by multiplication
print(f"\nx.is_leaf: {x.is_leaf}") # True — you created it
print(f"y.is_leaf: {y.is_leaf}") # False — PyTorch produced itVisualizing the Graph: Forward Pass Builds It, Backward Pass Walks It
The code above builds exactly this graph: z = x*y, then w = z + 4, then out = w**2. Toggle between the two views below — the forward pass shows data flowing left→right through each operation node; the backward pass shows .backward() walking the same graph in reverse, multiplying local derivatives via the chain rule to accumulate .grad at each leaf.
Forward pass for out = (x·y + 4)² at x=2, y=3: each circle is an operation PyTorch records (its grad_fn). Violet circles are leaf tensors you created; blue/green circles are non-leaf tensors PyTorch produced.
In TensorFlow 1.x, you had to define the graph explicitly before running it. In PyTorch, the graph is constructed automatically as Python code executes — this is the dynamic computation graph. Every time you do an operation on a tracked tensor, PyTorch adds a node to the graph. When you call .backward(), the graph is traversed in reverse to compute gradients. After the backward pass, the graph is freed from memory (to save RAM). This design is why you can use Python control flow (if/else, for loops) inside your model — the graph is different on every forward pass, and that's fine.
5 Automatic Differentiation: Your First Gradient
Let's compute gradients step-by-step and verify them against the analytical (hand-calculated) derivative. This demystifies what .backward() does.
import torch
# ── Example 1: Simple polynomial ──
# f(x) = x² + 2x + 1
# f'(x) = 2x + 2
# At x = 3: f'(3) = 2(3) + 2 = 8
x = torch.tensor(3.0, requires_grad=True)
y = x**2 + 2*x + 1
print(f"f(3) = 3² + 2(3) + 1 = {y.item()}") # 16.0
y.backward() # compute df/dx, store in x.grad
print(f"f'(3) = 2(3)+2 = {x.grad.item()}") # 8.0 ✓ matches analytical
print()
# ── Example 2: Chain of operations ──
# f(x) = (x² + 2x) * 3
# f'(x) = (2x + 2) * 3 = 6x + 6
# At x = 4: f'(4) = 6(4) + 6 = 30
x = torch.tensor(4.0, requires_grad=True)
inner = x**2 + 2*x # inner = x² + 2x
z = inner * 3 # z = 3(x² + 2x)
z.backward()
print(f"z at x=4: {z.item()}") # 3*(16+8) = 72.0
print(f"dz/dx at x=4: {x.grad.item()}") # 6*4+6 = 30.0 ✓
print()
# ── Example 3: Multi-variable gradients ──
# f(x, y) = x² * y + x * y²
# ∂f/∂x = 2xy + y²
# ∂f/∂y = x² + 2xy
# At x=2, y=3: ∂f/∂x = 12+9=21, ∂f/∂y = 4+12=16
x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)
f = x**2 * y + x * y**2
f.backward() # computes ∂f/∂x AND ∂f/∂y in one pass
print(f"f(2,3) = {f.item()}") # 4*3 + 2*9 = 12 + 18 = 30
print(f"∂f/∂x at (2,3): {x.grad}") # 21.0 ✓
print(f"∂f/∂y at (2,3): {y.grad}") # 16.0 ✓
# ── Example 4: Gradient of a vector output ──
# When output is a vector (not a scalar), need to specify gradient weights
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x**2 # y = [1, 4, 9] — vector output
# Calling y.backward() alone fails: "grad can be implicitly created only for scalar outputs"
# Solution 1: sum to get a scalar first
loss = y.sum() # loss = 1+4+9 = 14
loss.backward()
print(f"\nVector gradient (via .sum().backward()):")
print(f"d(sum(x²))/dx = 2x = {x.grad}") # [2., 4., 6.] ✓
# Solution 2: pass gradient tensor (for vector Jacobian products)
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x**2
grad_weights = torch.ones_like(y) # equivalent to .sum() trick
y.backward(grad_weights)
print(f"Vector backward with ones: {x.grad}") # [2., 4., 6.] ✓Gradient descent updates weights based on the gradient of a single scalar loss value: ∂L/∂w for each weight w. A loss function must return a scalar (a single number). When you call loss.backward(), PyTorch computes ∂L/∂w for every learnable parameter in the computation graph — all in one pass. If your output is a vector, you first reduce it to a scalar (usually by summing or taking the mean) to form a proper loss function. This is why the last line in every forward pass before .backward() always produces a scalar.
6 Gradient Accumulation and the .zero_grad() Pitfall
This is one of the most common bugs for PyTorch beginners. Understanding it will save you hours of debugging confusing training behavior.
PyTorch Accumulates Gradients by Default
When you call .backward(), PyTorch does not replace x.grad — it adds to it. This design choice exists because some training techniques (like gradient accumulation over multiple mini-batches to simulate larger batch sizes) need this behavior. But for standard training, you must manually zero the gradients before each backward pass.
import torch
# ── The Accumulation Bug ──
x = torch.tensor(2.0, requires_grad=True)
# Forward + backward pass 1
y = x ** 2 # dy/dx = 2x = 4
y.backward()
print(f"After 1st backward: x.grad = {x.grad}") # 4.0 ✓
# Forward + backward pass 2 — WITHOUT zeroing gradients
y = x ** 2 # same computation
y.backward()
print(f"After 2nd backward (no zero!): x.grad = {x.grad}") # 8.0 ✗ (should be 4.0!)
# Forward + backward pass 3 — correctly zeroed
x.grad.zero_() # in-place zero — ALWAYS do this before backward
y = x ** 2
y.backward()
print(f"After 3rd backward (with zero): x.grad = {x.grad}") # 4.0 ✓
print()
print("The rule: ALWAYS call zero_grad() before backward().")
print("In a model, this looks like: optimizer.zero_grad()")
# ── What the correct training loop looks like ──
# (Using a manual "optimizer" for illustration)
import torch.nn as nn
model = nn.Linear(3, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
X = torch.randn(10, 3)
y_true = torch.randn(10, 1)
for step in range(3):
# 1. Zero gradients FIRST (very important!)
optimizer.zero_grad()
# 2. Forward pass
y_pred = model(X)
loss = loss_fn(y_pred, y_true)
# 3. Backward pass (accumulate from zero, not from previous step)
loss.backward()
# 4. Update weights
optimizer.step()
print(f"Step {step+1}: loss = {loss.item():.4f}")Forgetting optimizer.zero_grad() causes gradients from the previous batch to add to the current batch's gradients. The symptom: loss decreases for a few steps then mysteriously blows up or oscillates wildly. The gradient updates become larger and larger each step. This is NOT caught by any warning — PyTorch happily accumulates gradients indefinitely. Make it a habit: the first line of every training loop iteration is always optimizer.zero_grad().
7 torch.no_grad(): Turning Off the Diary
The computation graph uses memory and computation resources. During inference (predicting on new data) or evaluation, you don't need gradients — you just want the forward pass output. torch.no_grad() disables gradient tracking for everything inside the block, making it faster and more memory-efficient.
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 10)
)
x = torch.randn(32, 784) # batch of 32 "images"
# ── Without no_grad: builds computation graph (wastes memory at inference) ──
y_with_grad = model(x)
print(f"With grad tracking: requires_grad={y_with_grad.requires_grad}") # True
# Memory: stores graph for all 32*784 → 256 → 10 computations
# ── With no_grad: no graph built ──
with torch.no_grad():
y_no_grad = model(x)
print(f"Without grad tracking: requires_grad={y_no_grad.requires_grad}") # False
# Memory: only stores the output values, not the graph
# ── As a decorator ──
@torch.no_grad()
def evaluate(model, X, y_true):
"""Always use no_grad for evaluation functions."""
logits = model(X)
predictions = logits.argmax(dim=1)
return predictions
# ── The model.eval() pattern ──
# model.eval() + torch.no_grad() are BOTH needed for correct evaluation:
# - model.eval() turns off Dropout and sets BatchNorm to use running stats
# - torch.no_grad() turns off gradient tracking
# Use them together EVERY time you evaluate:
model.eval()
with torch.no_grad():
outputs = model(x)
print(f"Evaluation output shape: {outputs.shape}") # (32, 10)
# After evaluation, return to training mode
model.train()
# ── requires_grad=False for frozen parameters (transfer learning preview) ──
# When fine-tuning, you often want to freeze early layers:
for param in model[0].parameters(): # freeze first Linear layer
param.requires_grad = False
# Check how many parameters are trainable
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
frozen = total - trainable
print(f"\nTotal params: {total:,}, Trainable: {trainable:,}, Frozen: {frozen:,}")8 Leaf Tensors vs Non-Leaf Tensors
PyTorch distinguishes between "leaf" tensors (tensors you create directly) and "non-leaf" tensors (tensors produced by operations). This distinction determines where gradients are stored — and understanding it helps you debug gradient issues.
import torch
# ── Leaf tensors: created directly by you ──
x = torch.tensor(2.0, requires_grad=True) # leaf
w = torch.tensor(3.0, requires_grad=True) # leaf
b = torch.tensor(1.0, requires_grad=False) # leaf (no grad tracking)
print(f"x.is_leaf: {x.is_leaf}") # True
print(f"w.is_leaf: {w.is_leaf}") # True
print(f"b.is_leaf: {b.is_leaf}") # True
# ── Non-leaf tensors: produced by operations ──
z = x * w + b # non-leaf: produced by multiplication and addition
output = z ** 2 # non-leaf: produced by squaring
print(f"\nz.is_leaf: {z.is_leaf}") # False
print(f"output.is_leaf: {output.is_leaf}") # False
# After backward, only leaf tensors with requires_grad=True store .grad
output.backward()
print(f"\nx.grad: {x.grad}") # ∂output/∂x = 2z * w = 2(x*w+b)*w = 2*(6+1)*3 = 42
print(f"w.grad: {w.grad}") # ∂output/∂w = 2z * x = 2*(6+1)*2 = 28
print(f"b.grad: {b.grad}") # None — b doesn't require grad
print(f"z.grad: {z.grad}") # None — non-leaf, grad not kept by default
# Verify manually
x_val, w_val, b_val = 2.0, 3.0, 1.0
z_val = x_val * w_val + b_val # = 7.0
# ∂output/∂x = d(z²)/dx = 2z * dz/dx = 2z * w = 2*7*3 = 42 ✓
# ∂output/∂w = d(z²)/dw = 2z * dz/dw = 2z * x = 2*7*2 = 28 ✓
# ── retain_grad(): keep gradients for non-leaf tensors (debugging) ──
x = torch.tensor(2.0, requires_grad=True)
z = x * 3
output = z ** 2
z.retain_grad() # "please keep my gradient even though I'm non-leaf"
output.backward()
print(f"\nWith retain_grad: z.grad = {z.grad}") # d(z²)/dz = 2z = 14.0
# Use retain_grad() when debugging: to check intermediate layer gradientsLeaf vs Non-Leaf: Where the Graph Boundary Lives
Every tensor you type into existence directly — inputs, weights, biases — is a leaf. Everything produced by an operation on tensors is non-leaf. Only leaves with requires_grad=True accumulate .grad; non-leaves are transient scaffolding that PyTorch frees after .backward() unless you call retain_grad().
After output.backward(), only the leaves x and w (both requires_grad=True) hold a populated .grad. b is a leaf too, but with requires_grad=False it stays None. The non-leaf tensors z and output never store .grad — that's exactly where model parameters (leaves) differ from activations (non-leaves) in a real network.
Model parameters (weights and biases in nn.Linear, nn.Conv2d, etc.) are leaf tensors — you create them when defining the model. Intermediate activations are non-leaf. During backpropagation, PyTorch only stores gradients in leaf tensors because those are the values you update with the optimizer. Storing gradients for all intermediate activations in a deep network would require enormous memory. If you need an intermediate gradient for debugging, use retain_grad(), but don't leave it on in production — it defeats the memory savings.
Real-World Spotlight: Gradient Descent from Scratch with Autograd
The best way to solidify your understanding of tensors and autograd is to implement gradient descent from scratch — using only tensor operations and .backward(), without any optimizer class. This is exactly what PyTorch does internally for you when you call optimizer.step().
import torch
# ── Problem: minimize a quadratic loss ──
# True parameters: w_true = 3.0, b_true = -2.0
# Loss: L(w, b) = (w - 3)² + (b + 2)²
# Minimum at w=3, b=-2 (both derivatives are zero there)
# Start from random initial values
torch.manual_seed(0)
w = torch.tensor(0.5, requires_grad=True) # starting guess
b = torch.tensor(1.0, requires_grad=True) # starting guess
lr = 0.1 # learning rate
print(f"Initial: w={w.item():.4f}, b={b.item():.4f}")
print(f"Target: w=3.0000, b=-2.0000")
print(f"\n{'Step':>5} | {'w':>8} | {'b':>8} | {'Loss':>10}")
print("-" * 42)
for step in range(25):
# Step 1: Zero gradients (IMPORTANT!)
if w.grad is not None: w.grad.zero_()
if b.grad is not None: b.grad.zero_()
# Step 2: Forward pass — compute loss
loss = (w - 3.0)**2 + (b + 2.0)**2
# Step 3: Backward pass — compute gradients
loss.backward()
# w.grad = ∂L/∂w = 2(w - 3)
# b.grad = ∂L/∂b = 2(b + 2)
# Step 4: Update parameters (gradient descent)
with torch.no_grad(): # don't track this update in the graph
w -= lr * w.grad
b -= lr * b.grad
if step % 5 == 0 or step == 24:
print(f"{step+1:>5} | {w.item():>8.4f} | {b.item():>8.4f} | {loss.item():>10.6f}")
print(f"\nFinal: w={w.item():.6f}, b={b.item():.6f}")
print(f"These should be very close to w=3.0, b=-2.0")Drag the slider to change the learning rate and watch how .backward()-computed gradients steer (w, b) across the loss surface L(w,b) = (w−3)² + (b+2)² toward the minimum at (3, −2), starting from the same (0.5, 1.0) as the code above:
lr = 0.10 — steady convergence toward (w=3, b=-2) in about 25 steps.
The parameters converge towards the true values (w=3, b=-2) through repeated gradient descent steps. This is gradient descent for a neural network — just with scalar parameters instead of millions. In a real neural network, the exact same process happens, but applied to millions of weights simultaneously. PyTorch's autograd handles all the derivative computation automatically through the chain rule, regardless of how complex the computation graph is.
Notice the with torch.no_grad(): wrapper around the weight update step. This is essential: the update w -= lr * w.grad is an arithmetic operation on a tracked tensor (w). Without no_grad, PyTorch would try to add this update operation to the computation graph, which would cause errors in the next backward pass. The update step is not part of the model's computation — it's the optimizer's job. Always wrap manual weight updates in torch.no_grad(). When using optimizer.step(), this is handled automatically.
✍️ Practice Exercises
- Tensor creation: Create the following tensors: (a) a 4×4 identity matrix as float32; (b) 100 evenly spaced values from 0 to 2π; (c) a batch of 8 "MNIST-like" images as random floats of shape (8, 1, 28, 28); (d) a tensor of ones matching the shape of a given tensor
t = torch.randn(3, 5). Print shapes and dtypes for all. - Matrix operations: Create two 3×3 matrices A and B. Compute: (a) A @ B (matrix multiply); (b) A * B (element-wise); (c) A.T (transpose); (d) the maximum value in each row of A; (e) the L2 norm (Euclidean length) of each row using
(A**2).sum(dim=1).sqrt(). - Gradient verification: For the function f(x) = sin(x) * exp(-x²/2) at x=1.0, compute f'(1) using PyTorch autograd. Then verify it analytically: f'(x) = cos(x)*exp(-x²/2) + sin(x)*(-x)*exp(-x²/2). Do they match?
- no_grad timing: Create a large model (
nn.Sequentialof 5 linear layers, 1000 neurons each). Time one forward pass with and withouttorch.no_grad(). How much faster is inference without gradient tracking? - Gradient descent: Extend the real-world example to minimize a three-variable quadratic: L(w1, w2, w3) = (w1-1)² + (w2-2)² + (w3-3)². The true minimum is (1, 2, 3). Run 50 steps and verify convergence.
▶ Show Key Solutions
import torch
import math
# Exercise 1
a = torch.eye(4) # (a)
b = torch.linspace(0, 2*math.pi, 100) # (b)
c = torch.rand(8, 1, 28, 28) # (c)
t = torch.randn(3, 5)
d = torch.ones_like(t) # (d)
for x, name in [(a,'eye'), (b,'linspace'), (c,'batch'), (d,'ones_like')]:
print(f"{name}: shape={x.shape}, dtype={x.dtype}")
# Exercise 3: Gradient verification
x = torch.tensor(1.0, requires_grad=True)
y = torch.sin(x) * torch.exp(-x**2 / 2)
y.backward()
print(f"\nAutograd f'(1) = {x.grad.item():.6f}")
x_val = 1.0
import math
analytical = (math.cos(x_val) * math.exp(-x_val**2/2)
+ math.sin(x_val) * (-x_val) * math.exp(-x_val**2/2))
print(f"Analytical f'(1) = {analytical:.6f}")
# Exercise 5: 3-variable gradient descent
w1 = torch.tensor(5.0, requires_grad=True)
w2 = torch.tensor(0.0, requires_grad=True)
w3 = torch.tensor(-3.0, requires_grad=True)
lr = 0.1
for step in range(50):
for p in [w1, w2, w3]:
if p.grad is not None: p.grad.zero_()
loss = (w1-1)**2 + (w2-2)**2 + (w3-3)**2
loss.backward()
with torch.no_grad():
w1 -= lr * w1.grad
w2 -= lr * w2.grad
w3 -= lr * w3.grad
print(f"\nFinal: w1={w1.item():.4f}, w2={w2.item():.4f}, w3={w3.item():.4f}")
# Should be close to (1.0, 2.0, 3.0)📚 Primary Sources for This Lesson
PyTorch Official: Automatic Differentiation with torch.autograd — the canonical tutorial from the PyTorch team. Short, clear, and authoritative. Covers the computation graph and gradient computation with simple examples.
PyTorch Official: Tensors Tutorial — comprehensive coverage of tensor operations, NumPy bridging, and device management.