Tensors and devices

A tensor is PyTorch's core data structure: a multi dimensional array, like a NumPy array, that can live on a GPU and track gradients.

Creating tensors

import torch

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
zeros = torch.zeros(3, 4)
rand = torch.randn(2, 3)          # normal distribution
x = torch.arange(10, dtype=torch.float32)

Every tensor has a shape, a dtype and a device:

rand.shape    # torch.Size([2, 3])
rand.dtype    # torch.float32
rand.device   # cpu

Most deep learning uses float32, and mixed precision training uses float16 or bfloat16 for speed. Mismatched dtypes are a common source of errors, so check them when something fails.

Operations and broadcasting

Arithmetic works element by element, and broadcasting stretches smaller tensors to match larger ones when shapes are compatible:

images = torch.randn(32, 3, 64, 64)   # batch, channels, height, width
mean = torch.tensor([0.5, 0.4, 0.3]).view(1, 3, 1, 1)
normalised = images - mean

The mean has shape (1, 3, 1, 1) and is broadcast across the batch and every pixel. Getting shapes right is a large part of deep learning programming. When in doubt, print .shape after every step.

Matrix multiplication uses @:

W = torch.randn(4, 3)
out = rand @ W.T   # (2, 3) @ (3, 4) gives (2, 4)

Devices

Training on a GPU can be dramatically faster. Pick a device once and move tensors and models to it:

if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")   # Apple Silicon
else:
    device = torch.device("cpu")

x = rand.to(device)

Operations between tensors on different devices raise an error. A typical bug is moving the model to the GPU but forgetting to move a batch of data.

Views and copies

Methods such as view, reshape, permute and slicing can return views that share memory with the original tensor. Changing a view changes the original. Use .clone() when you need an independent copy.

NumPy interoperability

torch.from_numpy(arr) and tensor.numpy() convert between libraries, sharing memory on the CPU. Tensors on a GPU must be moved to the CPU first with .cpu().

Practise by creating a batch of fake images, normalising it per channel and checking the resulting mean of each channel.