Tensores: do escalar ao lote de imagens

As dimensões, de zero a cinco

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

O último número é um aviso: um lote de 32 imagens de 224 por 224 em float32 ocupa cerca de 19 MB. Tensores crescem rápido, e o dtype e o tamanho do lote decidem se o seu computador aguenta.

Ler uma forma como um idioma

Quando um tutorial mostrar uma forma como (32, 224, 224, 3), leia como lote × altura × largura × canais. Esse hábito de leitura resolve a maior parte da confusão dos primeiros dias com aprendizado profundo.

Existem duas convenções para o lugar dos canais. O NumPy e o TensorFlow costumam usar canais no fim (altura, largura, canais), e o PyTorch usa canais antes (canais, altura, largura). O np.moveaxis troca de uma para a outra:

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

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

A mesma regra de sempre, em mais dimensões

Indexação, fatiamento, broadcasting e agregação funcionam exatamente como antes. Com axis você escolhe sobre quais dimensões resumir. Para calcular a média de cada canal de cor em um lote inteiro, resuma lote, altura e largura de uma vez, e sobram os canais:

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

(3,)

Uma imagem de verdade, em miniatura

Uma imagem colorida é um tensor. Em uma imagem de 2 por 2 pixels, cada pixel tem três canais (vermelho, verde, azul), de 0 a 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]]

Os quatro pixels são vermelho, verde, azul e branco. pixels[..., 0] extrai o canal vermelho, e a média dos três canais produz uma versão em tons de cinza (um jeito simples, e não o mais fiel à percepção humana, de converter).

O que o PyTorch e o JAX acrescentam

Os tensores do PyTorch e do JAX têm a mesma forma, a mesma indexação e as mesmas regras de broadcasting. O que eles acrescentam é execução em GPU e derivação automática, que calcula gradientes por você. Quem domina os arrays do NumPy já domina o vocabulário deles.