Tensors and dimensions
Translate an inspected NumPy model into tensors, modules, automatic differentiation and a clear training loop.
THE IDEA
A tensor is an array that knows its shape, type and device
If you have used NumPy, a tensor is the same idea: an n-dimensional array of numbers. It carries three pieces of metadata that matter constantly in practice: its shape (the size along each axis), its dtype (float or integer) and its device (CPU or a GPU). Most day-to-day bugs in deep learning code are one of these three not matching what you assumed.
The first axis of most data is the batch: many examples processed at once, so an image batch is (batch, channels, height, width) and a table batch is (batch, features). Operations are written to work on the whole batch in one go.
Predict first. Is * the same operation as @ for two matrices?
Play with shapes
The shape rules are small but unforgiving. reshape may only rearrange, so the element count must stay the same, and a -1 lets PyTorch work out one size. unsqueeze inserts a size-1 axis. Broadcasting lines shapes up from the right and stretches any axis of size 1.
Predict first. You add a tensor of shape (8, 10) and one of shape (8,). What happens?
import torch
x = torch.zeros(8, 28, 28)
x.reshape(8, -1).shape # torch.Size([8, 784])
x.unsqueeze(1).shape # torch.Size([8, 1, 28, 28])
(torch.zeros(8, 10) + torch.zeros(10)).shape # torch.Size([8, 10])
(torch.zeros(32, 10) @ torch.zeros(10, 3)).shape # torch.Size([32, 3])Chapter context
PyTorch provides tensor operations and differentiation infrastructure. The framework removes bookkeeping, but it does not remove the need to understand dimensions, parameter ownership, device placement and evaluation behavior.
TRY IT YOURSELF
Part 1. Do it hands-on
PyTorch itself cannot run in a browser here, so the Python parts use a small PyTorch-style library that is already loaded. It has the same names and call shapes (torch.tensor, nn.Linear, loss.backward(), optimizer.step(), DataLoader) and the same behaviour for everything these lessons teach: shapes, accumulating .grad, graphs freed after backward(), parameter registration. It is not PyTorch: initial values and random numbers differ, and "cuda" is only a label. Every code block marked real PyTorch is for you to run on your own machine.
First predict shapes by hand. Type each result like (8, 784), or the word error when PyTorch would refuse.
| setup | expression | resulting shape |
|---|---|---|
| x = torch.zeros(8, 28, 28) | x.reshape(8, -1) | |
| x = torch.zeros(8, 28, 28) | x.unsqueeze(1) | |
| x = torch.zeros(8, 28, 28) | x.transpose(1, 2) | |
| a = torch.zeros(8, 10); b = torch.zeros(10) | a + b | |
| a = torch.zeros(8, 10); b = torch.zeros(8) | a + b | |
| a = torch.zeros(8, 1); b = torch.zeros(1, 10) | a + b | |
| a = torch.zeros(32, 10); b = torch.zeros(10, 3) | a @ b | |
| a = torch.zeros(32, 10); b = torch.zeros(5, 3) | a @ b | |
| m = torch.zeros(4, 3) | m.sum(dim=1, keepdim=True) |
- To do: Reshaping: predict the result of the three operations on x of shape (8, 28, 28). Type each shape like (8, 784)
- To do: Broadcasting: predict a + b for the three pairs. Type error when PyTorch cannot match them
- To do: Matrix products and sums: predict a @ b (one of them fails) and m.sum(dim=1, keepdim=True)
Part 2. Do it in code
- torch and nn are available (a small PyTorch-style library, see the note above). Write the four functions below using tensor operations.
- flatten_images(x): x has shape (B, ...). Return it as (B, N) with reshape, keeping the batch axis.
- add_channel_dim(x): (B, H, W) -> (B, 1, H, W) with unsqueeze.
- center_features(x): x is (B, F). Subtract each feature's mean over the batch (mean over dim 0) so every column averages 0.
- pairwise_scores(a, b): a is (n, d) and b is (m, d). Return the (n, m) matrix of dot products a @ b.T.
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Is * the same operation as @ for two matrices?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
For a batch X of shape (B,d), a dense output with h features has shape (B,h). torch.nn.Linear stores weight with shape (h,d) and computes X times weight-transpose plus bias. Parameter count is d*h+h. PyTorch execution belongs in a linked notebook or supported runtime; the chapter does not assume browser PyTorch.
Worked example
A batch has shape (2,4). A Linear(4,3) layer stores a (3,4) weight and a length-three bias. Its output is (2,3). The parameter count is 4×3+3=15. This differs from the batch's element count: model parameters are shared across both examples.
Reference code
The following example uses the Python standard library and does not call a remote service.
batch, inputs, outputs = 2,4,3
weight_shape = (outputs,inputs)
output_shape = (batch,outputs)
assert weight_shape == (3,4)
assert output_shape == (2,3)
assert inputs*outputs+outputs == 15
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
eval mode and no_grad are independent. Model parameters do not multiply with batch size. A tensor detached for logging must not replace the tensor used to compute the training loss.