PyTorch Tensors and Autograd for Beginners

PyTorch is the most widely used framework for deep learning in research and industry. It gives you two things: tensors (like NumPy arrays, but able to run on GPUs) and autograd (automatic calculation of gradients, so you never write backpropagation by hand).

Installing#

pip install torch

The PyTorch website gives the exact command for your system, including GPU (CUDA) versions. The examples here were run with the CPU version, torch 2.14.

Creating tensors#

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(a)
print("shape:", a.shape, " dtype:", a.dtype)

print(torch.zeros(2, 3))
print(torch.arange(0, 10, 3))
f = torch.tensor([1.5, 2.5])
print(f.dtype)

tensor([[1, 2, 3],
        [4, 5, 6]])
shape: torch.Size([2, 3])  dtype: torch.int64
tensor([[0., 0., 0.],
        [0., 0., 0.]])
tensor([0, 3, 6, 9])
torch.float32

Integers default to int64 and decimals to float32. Neural networks work in float32, so convert with .float() when needed.

Maths, shapes and broadcasting#

Tensor operations look just like NumPy:

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(x * 10)
print(x @ x)                       # matrix multiplication
print(x.sum(), x.mean(dim=0))      # dim works like NumPy's axis
print(x + torch.tensor([100.0, 200.0]))   # broadcasting

images = torch.zeros(32, 1, 28, 28)        # a batch: 32 images, 1 channel, 28x28 pixels
print(images.shape, "→", images.view(32, -1).shape)

tensor([[10., 20.],
        [30., 40.]])
tensor([[ 7., 10.],
        [15., 22.]])
tensor(10.) tensor([2., 3.])
tensor([[101., 202.],
        [103., 204.]])
torch.Size([32, 1, 28, 28]) → torch.Size([32, 784])

.view (or .reshape) changes the shape without copying data; -1 means "work out this size". Image batches in PyTorch are shaped (batch, channels, height, width).

NumPy and PyTorch together#

import numpy as np
import torch

arr = np.array([1.0, 2.0, 3.0])
t = torch.from_numpy(arr)
print(t, t.dtype)
print(t.float().numpy())
print(torch.tensor([[0.25]]).item())       # a single value as a Python number

tensor([1., 2., 3.], dtype=torch.float64) torch.float64
[1. 2. 3.]
0.25

torch.from_numpy shares memory with the array (changing one changes the other); torch.tensor(arr) makes a copy.

Devices: CPU and GPU#

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
print("using:", device)
x = torch.ones(3).to(device)
print(x.device.type)

using: cpu
cpu

Every tensor and the model must be on the same device. This example ran on a machine without a GPU, so it uses the CPU; the same code uses a GPU automatically where one exists. Apple Silicon Macs use "mps".

Autograd: automatic gradients#

Mark a tensor with requires_grad=True, compute something from it, and call .backward(). PyTorch fills in .grad with the derivative:

import torch

w = torch.tensor(3.0, requires_grad=True)
loss = (w - 5) ** 2              # smallest when w = 5
loss.backward()
print("loss:", loss.item(), " gradient:", w.grad.item())

loss: 4.0  gradient: -4.0

The derivative of (w − 5)² is 2(w − 5), which is −4 at w = 3. The negative gradient says "increase w to reduce the loss".

Gradient descent by hand#

Fit a line y = w·x + b with autograd doing the calculus:

import torch

torch.manual_seed(0)
x = torch.linspace(0, 10, 50)
y = 3 * x + 2 + torch.randn(50) * 0.5          # true slope 3, intercept 2

w = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.01
for step in range(1, 2001):
    loss = ((w * x + b - y) ** 2).mean()
    loss.backward()
    with torch.no_grad():                      # update without tracking gradients
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad.zero_(); b.grad.zero_()             # gradients add up unless cleared
    if step in (1, 500, 2000):
        print(f"step {step:>4}: loss {loss.item():.3f}, w {w.item():.3f}, b {b.item():.3f}")

step    1: loss 369.674, w 2.228, b 0.340
step  500: loss 0.286, w 3.045, b 1.769
step 2000: loss 0.282, w 3.029, b 1.881

The parameters approach the true slope 3 and intercept 2. Three details matter in every PyTorch training loop:

• update weights inside torch.no_grad();
• clear gradients each step, because .backward() adds to .grad;
• the loss must be a single number to call .backward() on it.

The next lesson replaces the manual updates with PyTorch's nn modules and optimisers.

Common mistakes#

Practice#

import torch

x = torch.tensor(2.0, requires_grad=True)
f = x ** 3 + 2 * x
f.backward()
print("autograd:", x.grad.item(), " by hand:", 3 * 2 ** 2 + 2)

autograd: 14.0  by hand: 14

Key takeaway#

Next: Training a Neural Network in PyTorch.