Tensors & autograd

Everything in PyTorch — every image, every network weight, every prediction — is a tensor. A tensor is just an n-dimensional grid of numbers: a single number is a 0-D tensor, a list is 1-D, a table is 2-D, a stack of tables is 3-D, and so on. If you've used NumPy arrays, you already know the idea. PyTorch adds two superpowers on top: tensors can run on a GPU, and they can track their own gradients.

You can follow along in a Python shell (python) or a scratch file. First, make a few:

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])   # a 2×2 tensor
print(x)
print("shape:", x.shape)     # torch.Size([2, 2])
print("mean:", x.mean())     # tensor(2.5000)

zeros = torch.zeros(2, 3)    # a 2×3 grid of zeros
rand  = torch.rand(2, 3)     # random numbers in [0, 1)

Shapes are everything

Most PyTorch bugs are shape mismatches, so it pays to think in shapes. Two moves you'll use constantly are reshaping and matrix multiplication:

x = torch.rand(4, 28, 28)    # 4 images, each 28×28 pixels
flat = x.reshape(4, 28 * 28) # flatten each image to a 784-long row
print(flat.shape)            # torch.Size([4, 784])

a = torch.rand(2, 3)
b = torch.rand(3, 5)
print((a @ b).shape)         # torch.Size([2, 5]) — matrix multiply

Our digit images will arrive as 28×28 grids; the network wants each one flattened into a row of 784 numbers. That reshape is exactly the move we'll make in Chapter 4.

Moving to your device

Remember the device from the last chapter? Sending a tensor to it is one call. For a network to train, its data and its weights must live on the same device:

device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.rand(2, 3).to(device)
print(x.device)

Autograd: the reason PyTorch exists

Here's the magic that makes training possible. Tell a tensor to track gradients with requires_grad=True, and PyTorch quietly records every operation you do to it. Call .backward() on the result and it computes the derivative — the gradient — of that result with respect to your tensor, automatically:

w = torch.tensor(3.0, requires_grad=True)
y = w ** 2 + 2 * w     # some function of w

y.backward()           # compute dy/dw
print(w.grad)          # tensor(8.) — because dy/dw = 2w + 2 = 8 at w=3

That w.grad is the number that says "nudge w this way to change y." Training a neural network is nothing more than doing this for millions of weights at once and taking a tiny step in the direction that lowers the error — a topic the backpropagation course covers in depth. You won't call .backward() by hand in this project; PyTorch does it inside the training loop. But now you know what it's doing. Time to get some real data.