PyTorch

Tenseurs et device

import torch

x = torch.tensor([[1., 2.], [3., 4.]])
torch.zeros(2, 3); torch.ones(2, 3); torch.randn(2, 3)
torch.arange(0, 10, 2); torch.linspace(0, 1, 5)

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
print(device, torch.__version__)

mps 2.9.1

x = x.to(device)
model = model.to(device)

Erreur la plus fréquente : Expected all tensors to be on the same device. Les données et le modèle doivent être sur le même device, à chaque batch.

Formes

x.shape, x.dtype, x.device
x.view(-1, 4)         # nécessite un tenseur contigu
x.reshape(-1, 4)      # marche toujours, copie si besoin
x.permute(0, 2, 1)    # réordonne les axes
x.unsqueeze(0)        # ajoute un axe -> (1, ...)
x.squeeze()           # retire les axes de taille 1
torch.cat([a, b], dim=0)
torch.stack([a, b])   # crée un nouvel axe

torch.einsum("bij,bjk->bik", a, b) remplace avantageusement les enchaînements de permute + matmul quand la manipulation d'axes devient illisible.

Autograd

w = torch.randn(3, requires_grad=True)
loss = (w ** 2).sum()
loss.backward()
print(w.grad)

tensor([ 1.4832, -0.6210,  2.0044])

with torch.no_grad():        # désactive le graphe : inférence, évaluation
    preds = model(x)

x.detach()                   # coupe un tenseur du graphe

Un modèle

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, d_in: int, d_hidden: int, d_out: int):
        super().__init__()                    # obligatoire, avant tout le reste
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(d_hidden, d_out),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = MLP(128, 512, 10).to(device)
print(sum(p.numel() for p in model.parameters()) / 1e6, "M paramètres")

0.13 M paramètres

Boucle d'entraînement

from torch.utils.data import DataLoader, TensorDataset

loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True, num_workers=4)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
criterion = nn.CrossEntropyLoss()

for epoch in range(epochs):
    model.train()
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)

        opt.zero_grad(set_to_none=True)     # sinon les gradients s'accumulent
        loss = criterion(model(xb), yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

    sched.step()

    model.eval()
    with torch.no_grad():
        acc = (model(X_val.to(device)).argmax(1) == y_val.to(device)).float().mean()
    print(f"epoch {epoch}  loss {loss.item():.4f}  val_acc {acc:.3f}")

epoch 0  loss 1.8342  val_acc 0.412
epoch 1  loss 1.1907  val_acc 0.638
epoch 2  loss 0.8455  val_acc 0.741

Les quatre oublis classiques : zero_grad absent, model.train() / model.eval() non basculés (dropout et batchnorm se comportent différemment), no_grad manquant en évaluation, et données restées sur le CPU.

Pertes

Tâche	Perte	Entrée attendue
Classification multi-classe	nn.CrossEntropyLoss	logits bruts, pas de softmax
Classification binaire	nn.BCEWithLogitsLoss	logits bruts
Régression	nn.MSELoss, nn.L1Loss	valeurs

CrossEntropyLoss applique le log-softmax en interne. Y ajouter un softmax dans le modèle est une erreur silencieuse qui dégrade l'apprentissage sans lever.

Précision mixte

scaler = torch.amp.GradScaler(device)

with torch.autocast(device_type=device, dtype=torch.bfloat16):
    loss = criterion(model(xb), yb)

scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()

En bf16 le GradScaler est facultatif ; il reste nécessaire en fp16.

Sauvegarder

torch.save(model.state_dict(), "model.pt")
model.load_state_dict(torch.load("model.pt", map_location=device))
model.eval()

Toujours sauvegarder le state_dict, jamais l'objet modèle : la sérialisation directe casse dès que le code de la classe change.

Checkpoint complet pour reprendre un entraînement :

torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch}, "ckpt.pt")

Diagnostic mémoire GPU

print(torch.cuda.memory_allocated() / 1e9, "Go")
torch.cuda.empty_cache()

CUDA out of memory : réduire le batch, activer l'accumulation de gradients, passer en précision mixte, ou activer le gradient checkpointing.

Accélérer

model = torch.compile(model)       # gains réels sur GPU récents
torch.backends.cuda.matmul.allow_tf32 = True

See also

• numpy
• keras-tensorflow