M2.3 — PyTorch Tensors and Tensor Operations

Learning objectives

After this lesson, you should be able to:

• Create tensors and read their shape, dtype and device.
• Apply NumPy-style indexing, broadcasting, reductions and matrix multiplication to tensors.
• Convert between NumPy arrays and tensors and know they share memory.
• Arrange data as a batch, and image data channels first.
• Compute a gradient with requires_grad and backward().
• Explain torch.nn.Linear as Wx + b and read its parameter shapes.

The one-sentence idea

A tensor is an array that knows its shape, dtype and device, and, when asked, how it was computed.

Same ideas as NumPy, plus a GPU and gradients.

What a tensor is, and why PyTorch

A tensor is PyTorch’s array: a grid of numbers with a shape, a dtype and a location (CPU or GPU). If you read M2.2, almost everything here will look familiar, because PyTorch deliberately mirrors NumPy’s vocabulary. What a tensor adds is two things that training needs: it can live on a GPU, and PyTorch can track the operations applied to it so it can compute gradients automatically.

Most open robot-learning code, including the training code behind VLA models, is written with PyTorch. Being comfortable with tensors is therefore a direct prerequisite for reading it.

Creating tensors and reading their properties

A tensor has a shape, a dtype and a device. As with arrays, checking these is the first step when code misbehaves.

import torch

t = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])
print(t.shape, t.dtype, t.device)
print(torch.zeros(2, 3))

torch.Size([2, 3]) torch.float32 cpu
tensor([[0., 0., 0.],
        [0., 0., 0.]])

Notice the default float type is float32 (32 bits), not NumPy’s float64. Neural networks normally use 32-bit or smaller floats because they are faster and use half the memory, and the extra precision is not needed. Integer lists produce int64 tensors, so mixing integers and floats sometimes needs an explicit conversion:

i = torch.tensor([1, 2, 3])
print(i.dtype)
print(i.float().dtype)
print(i / 2)

torch.int64
torch.float32
tensor([0.5000, 1.0000, 1.5000])

Indexing, arithmetic and matrix multiplication

Indexing, slicing, masks, broadcasting and reductions work as in NumPy, and @ is matrix multiplication exactly as in M1.1. Reductions take dim instead of axis (PyTorch accepts both spellings in many places).

print(t[:, 1])                          # the second column
print(t + torch.tensor([10., 20., 30.]))  # broadcasting, shape (2, 3) + (3,)
print(t @ t.T)                          # (2, 3) times (3, 2) gives (2, 2)
print(t.mean(dim=0))                    # one mean per column

tensor([2., 5.])
tensor([[11., 22., 33.],
        [14., 25., 36.]])
tensor([[14., 32.],
        [32., 77.]])
tensor([2.5000, 3.5000, 4.5000])

The shape rule from M1.1 is unchanged: (2, 3) @ (3, 2) is valid and gives (2, 2). When a PyTorch script fails with a shape error, writing the shapes down and checking the inner dimensions finds the cause most of the time.

Moving between NumPy and PyTorch

torch.from_numpy(a) makes a tensor from an array, and t.numpy() goes the other way (for a CPU tensor). Both share memory with the original, like the views in M2.2, so changing one changes the other.

import numpy as np

a = np.arange(3.0)
s = torch.from_numpy(a)
s[0] = 5.0
print(a)

[5. 1. 2.]

Data from a camera or a dataset usually arrives as NumPy arrays and is converted to tensors at the start of the training or inference code.

Batches and image layout

Models process several examples at once, so tensors normally have a leading batch dimension. A batch of 8 vectors of length 5 has shape (8, 5). A batch of images has four dimensions, and here PyTorch differs from the NumPy convention in M2.2: images are stored channels first, as (batch, channels, height, width), not (height, width, channels).

img = torch.zeros(4, 6, 3)          # height 4, width 6, 3 colour channels (NumPy-style)
chw = img.permute(2, 0, 1)          # reorder the dimensions: channels first
batch = chw.unsqueeze(0)            # add a leading batch dimension
print(chw.shape)
print(batch.shape)

torch.Size([3, 4, 6])
torch.Size([1, 3, 4, 6])

permute reorders dimensions, unsqueeze adds a size-1 dimension, and reshape changes the shape while keeping the number of entries. Forgetting the channel order is a classic source of silent errors, because a wrongly ordered image still has a valid shape for some operations.

Automatic differentiation

If a tensor is created with requires_grad=True, PyTorch records the operations performed on it. Calling .backward() on a result then computes the gradient of that result with respect to the tensor and stores it in .grad. This is the mechanism that trains every network: you will use it, not implement it, until the backpropagation lesson explains what happens inside.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x          # y = x^2 + 3x
y.backward()                # compute dy/dx
print(x.grad)               # the derivative 2x + 3 at x = 2

tensor(7.)

The derivative of x² + 3x is 2x + 3, which is 7 at x = 2, and PyTorch found it without being given a formula. The calculus behind this is the subject of M1.2.

A linear layer in PyTorch

The y = Wx + b layer from M1.1 is torch.nn.Linear. It holds a weight tensor and a bias tensor, which are the learned parameters, and calling the layer applies it. Following the note in M1.1, PyTorch stores the weight as (out, in).

torch.manual_seed(0)                       # make the random initial weights repeatable
layer = torch.nn.Linear(3, 2)              # 3 inputs, 2 outputs
x = torch.ones(3)
print(layer.weight.shape, layer.bias.shape, layer(x).shape)
print(torch.allclose(layer(x), x @ layer.weight.T + layer.bias))

torch.Size([2, 3]) torch.Size([2]) torch.Size([2])
True

Because the layer is an object that can be called, this is the class pattern from M2.1 (__init__ to set up, __call__ to run). Larger models are classes that hold several such layers.

Running on a GPU

A GPU performs the large array operations of training far faster than a CPU. A tensor lives on one device, and an operation needs all of its tensors on the same one. You move a tensor with .to(device):

device = "cuda" if torch.cuda.is_available() else "cpu"
t = t.to(device)
print(t.device)

The output depends on the machine, so it is not shown here. If you combine a CPU tensor and a GPU tensor you get an error that mentions two devices, and the fix is to move one of them. M2.5 covers using the GPU inside a training loop. Do not worry about having a GPU for this prerequisite material, since every example above runs on a CPU.

Common misconceptions

Engineering takeaways

1. A tensor has a shape, a dtype and a device. Print them when debugging.
2. Indexing, broadcasting and reductions work as in NumPy, with @ for matrix multiplication.
3. The default float type is float32, and integer lists give int64.
4. from_numpy and .numpy() share memory.
5. Batches add a leading dimension, and images are normally (batch, channels, height, width).
6. requires_grad=True and backward() compute gradients automatically.
7. torch.nn.Linear stores a weight of shape (out, in) and a bias, and computes Wx + b.
8. All tensors in an operation must be on the same device.

Check the shape, the dtype and the device, and most tensor errors explain themselves.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

You can now work with tensors and have seen autograd and nn.Linear. Next is how a network of such layers learns from its mistakes.

Back to the VLA roadmap prerequisites