Tensors: from scalar to batch of images

The dimensions, from zero to five

import numpy as np

escalar = np.array(7)
vetor = np.array([255, 0, 0])
matriz = np.zeros((28, 28))
imagem = np.zeros((224, 224, 3), dtype=np.uint8)
lote = np.zeros((32, 224, 224, 3), dtype=np.float32)

for nome, t in [("escalar", escalar), ("vetor", vetor), ("matriz", matriz), ("imagem", imagem), ("lote", lote)]:
    print(f"{nome:<8} ndim={t.ndim} shape={t.shape}")
print(round(lote.nbytes / 1e6, 1), "MB")

escalar  ndim=0 shape=()
vetor    ndim=1 shape=(3,)
matriz   ndim=2 shape=(28, 28)
imagem   ndim=3 shape=(224, 224, 3)
lote     ndim=4 shape=(32, 224, 224, 3)
19.3 MB

The last number is a warning: a batch of 32 images of 224 by 224 in float32 takes about 19 MB. Tensors grow quickly, and the dtype and the batch size determine if your computer can handle it.

Read a shape like a language

When a tutorial shows a shape like (32, 224, 224, 3), read as batch × height × width × channels. This reading habit resolves most of the confusion in the early days of deep learning.

There are two conventions for the position of the channels. NumPy and TensorFlow usually use channels at the end (altura, largura, canais), and PyTorch uses channels at the beginning (canais, altura, largura). The np.moveaxis switch from one to the other:

chw = np.moveaxis(imagem, -1, 0)
print(chw.shape, np.transpose(chw, (1, 2, 0)).shape)

(3, 224, 224) (224, 224, 3)

The same rule as always, in more dimensions

Indexing, slicing, broadcasting, and aggregation work exactly as before. With axis you choose about which dimensions to summarize. To calculate the average of each color channel in an entire batch, summarize batch, height, and width at once, and the channels remain:

print(lote.mean(axis=(0, 1, 2)).shape)

(3,)

A real image, in miniature

A color image is a tensor. In a 2 by 2 pixel image, each pixel has three channels (red, green, blue), from 0 to 255:

pixels = np.array(
    [[[255, 0, 0], [0, 255, 0]], [[0, 0, 255], [255, 255, 255]]],
    dtype=np.uint8,
)
print(pixels.shape)
print(pixels[0, 1])
print(pixels[..., 0])
cinza = pixels.mean(axis=-1).astype(np.uint8)
print(cinza)

(2, 2, 3)
[  0 255   0]
[[255   0]
 [  0 255]]
[[ 85  85]
 [ 85 255]]

The four pixels are red, green, blue, and white. pixels[..., 0] extracts the red channel, and the average of the three channels produces a grayscale version (a simple way, and not the most faithful to human perception, to convert).

What PyTorch and JAX add

The tensors of PyTorch and JAX have the same shape, the same indexing, and the same broadcasting rules. What they add is execution on GPU and automatic differentiation, which calculates gradients for you. Anyone who masters NumPy arrays already masters their vocabulary.