PyTorch and Autograd¶
import torch
import numpy as np
1. Tensors are essentially NumPy arrays that can track gradients¶
Most NumPy habits carry over directly: indexing, slicing, broadcasting, @ for matrix multiplication.
A = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
v = torch.tensor([1.0, 1.0])
print(A @ v)
print(A * 10 + 1)
print(A.shape, A.dtype)
tensor([3., 7.])
tensor([[11., 21.],
[31., 41.]])
torch.Size([2, 2]) torch.float32
Gotcha: NumPy defaults to float64, PyTorch to float32, and converting a NumPy array keeps its type. Watch what happens when that meets a float32 weight matrix:
X_np = np.random.random((5, 2)) # NumPy: float64
X = torch.tensor(X_np) # still float64
W = torch.randn(2, 1) # PyTorch: float32
print(X.dtype, W.dtype)
Y = X @ W # Oh no!
torch.float64 torch.float32
--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) Cell In[3], line 6 3 W = torch.randn(2, 1) # PyTorch: float32 5 print(X.dtype, W.dtype) ----> 6 Y = X @ W # Oh no! RuntimeError: expected m1 and m2 to have the same dtype, but got: double != float
X = torch.tensor(X_np, dtype=torch.float32)
Y = X @ W
print(Y.dtype)
torch.float32
2. Autograd: forward builds a graph, backward() fills in .grad¶
Mark the tensors you want derivatives with respect to using requires_grad=True. Every operation on them is recorded.
x = torch.tensor(3.0, requires_grad=True)
y = torch.tensor(4.0, requires_grad=True)
z = x * y + x**2
print(z) # note the grad_fn: z remembers how it was computed, take a look at .next_functions
tensor(21., grad_fn=<AddBackward0>)
z.backward() # walk the graph backward (reverse-mode autodiff)
print(x.grad) # dz/dx = y + 2x = 10
print(y.grad) # dz/dy = x = 3
tensor(10.) tensor(3.)
3. Gradients accumulate¶
backward() adds to .grad; it does not overwrite it. Every training loop needs to reset the gradients each step.
w = torch.tensor(3.0, requires_grad=True)
for step in range(3):
loss = w**2
loss.backward()
print(w.grad) # 6, then 12, then 18: wrong after the first step!
tensor(6.) tensor(12.) tensor(18.)
w.grad.zero_() # trailing underscore = in-place operation
print(w.grad)
tensor(0.)
4. How to write the update step¶
You want $w \leftarrow w - \eta \nabla_w L$. There is exactly one idiom you should use:
with torch.no_grad():
w -= learning_rate * w.grad
torch.no_grad(): the update is not part of the model, so don't record it in the graph.-=: update the existing tensor in place.
Here are two ways to get it wrong:
# Wrong #1: forget no_grad
w = torch.tensor([1.0, 2.0], requires_grad=True)
loss = (w**2).sum()
loss.backward()
w -= 0.1 * w.grad # Bad weight update!
--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) Cell In[9], line 6 3 loss = (w**2).sum() 4 loss.backward() ----> 6 w -= 0.1 * w.grad # Bad weight update! RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.
# Wrong #2: use no_grad, but don't update in place
w = torch.tensor([1.0, 2.0], requires_grad=True)
loss = (w**2).sum()
loss.backward()
with torch.no_grad():
w = w - 0.1 * w.grad # creates a NEW tensor named w
print(w.requires_grad) # False: w is no longer a parameter
print(w.grad) # None: the next backward() will fail
False None
# Correct!
w = torch.tensor([1.0, 2.0], requires_grad=True)
loss = (w**2).sum()
loss.backward()
with torch.no_grad():
w -= 0.1 * w.grad
w.grad.zero_()
print(w, w.requires_grad)
tensor([0.8000, 1.6000], requires_grad=True) True
5. Store losses with .item()¶
loss is a tensor attached to the whole computation graph. .item() gives you a plain Python number.
losses = []
loss = (w**2).sum()
print(loss) # tensor with a grad_fn
print(loss.item()) # just a float
losses.append(loss.item()) # do this, not losses.append(loss)
tensor(3.2000, grad_fn=<SumBackward0>) 3.200000286102295
6. Shapes: one sample vs. a batch¶
The same expression works on a single sample or on a whole batch, because @ and + broadcast.
X = torch.rand(100, 2) # 100 samples, 2 features
w = torch.randn(2, 1)
b = torch.randn(1)
print(X[0, :].shape) # one sample: (2,)
print(X[0:1, :].shape) # slice keeps the batch dimension: (1, 2)
print((X[0, :] @ w + b).shape) # (1,)
print((X @ w + b).shape) # (100, 1): all samples at once, '+ b' is broadcast
torch.Size([2]) torch.Size([1, 2]) torch.Size([1]) torch.Size([100, 1])
Today you'll train one sample at a time, but next time we'll process whole batches. Write your forward function so it works on both.
8. Where this is going¶
Today we build everything by hand. Soon you'll use PyTorch's building blocks, which do exactly the same things:
| Today (by hand) | PyTorch library |
|---|---|
W1, b1 tensors with requires_grad=True |
nn.Linear(2, 20) |
relu(...) |
nn.ReLU() |
logistic(...) + cross-entropy |
nn.BCEWithLogitsLoss() |
your sgd(params) function |
torch.optim.SGD(params, lr).step() |
param.grad.zero_() |
optimizer.zero_grad() |
Under the hood, optimizer.step() is a torch.no_grad() block with an in-place update, just like the one you will write.
import torch.nn as nn
model = nn.Sequential(nn.Linear(2, 20), nn.ReLU(), nn.Linear(20, 1))
print(model)
for p in model.parameters():
print(p.shape, p.requires_grad)
Sequential( (0): Linear(in_features=2, out_features=20, bias=True) (1): ReLU() (2): Linear(in_features=20, out_features=1, bias=True) ) torch.Size([20, 2]) True torch.Size([20]) True torch.Size([1, 20]) True torch.Size([1]) True
Notice that nn.Linear stores its weight matrix as (outputs, inputs), the transpose of the W1 you will create today. It computes x @ W.T + b, so the math is the same.
The whole training loop, using the library pieces. Note that the model has no logistic layer at the end: BCEWithLogitsLoss applies it inside the loss.
X = torch.rand(100, 2)
y = ((X[:, 0] > 0.5) != (X[:, 1] > 0.5)).float().unsqueeze(1) # XOR labels
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for epoch in range(100):
total_loss = 0
for i in range(X.shape[0]):
loss = loss_fn(model(X[i:i+1]), y[i:i+1]) # a "batch" of one: shape (1, 2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if epoch % 20 == 0:
print(f"epoch {epoch:3d} loss {total_loss:.2f}")
with torch.no_grad():
pred = torch.sigmoid(model(X)) > 0.5
print(f"accuracy: {(pred == y).float().mean().item():.2f}")
epoch 0 loss 72.98 epoch 20 loss 21.74
epoch 40 loss 17.62 epoch 60 loss 15.79
epoch 80 loss 14.97 accuracy: 0.91