PyTorch : définition, fonctionnement et exemple

PyTorch est la bibliothèque la plus utilisée pour construire et entraîner des réseaux de neurones en Python. Elle fournit des tenseurs calculables sur GPU, un moteur de différentiation automatique et une bibliothèque de couches, de fonctions de perte et d'optimiseurs. Elle constitue la base pratique du Deep Learning moderne.

PyTorch se situe sous les outils de plus haut niveau : les Transformers de Hugging Face, les moteurs comme vLLM ou les bibliothèques de fine-tuning s'appuient dessus. Il appartient à la famille du Machine Learning, mais il vise les modèles de grande taille entraînés par descente de gradient, là où des bibliothèques comme scikit-learn traitent les modèles classiques.

Le projet est né à Facebook AI Research (aujourd'hui Meta AI), autour d'Adam Paszke, Sam Gross, Soumith Chintala et Gregory Chanan, et s'inspire de la bibliothèque Torch, écrite en Lua. Il est publié en open source en 2016 et présenté dans un article à NeurIPS en 2019. En 2022, sa gouvernance est confiée à la PyTorch Foundation, rattachée à la Linux Foundation, ce qui l'a placé sous une direction neutre vis-à-vis de Meta. La version 2.0, publiée en mars 2023, a introduit torch.compile.

À retenir : PyTorch n'est pas un modèle ni une plateforme, c'est une bibliothèque de calcul. Il fournit les briques pour écrire, entraîner et exécuter un modèle, mais ne décide ni de l'architecture ni des données.

Trois mécanismes expliquent la façon dont PyTorch fonctionne et pourquoi il s'est imposé auprès des chercheurs.

Les tenseurs

Un tenseur est un tableau multidimensionnel : un nombre, un vecteur, une matrice ou un volume plus complexe, comme un lot d'images. Il possède un type de données (float32, bfloat16...) et un périphérique, CPU ou GPU. L'interface ressemble à celle de NumPy, avec une différence décisive : un tenseur peut être déplacé sur un GPU, où les opérations matricielles s'exécutent en parallèle. L'accélération s'appuie sur la plateforme CUDA de NVIDIA, avec une prise en charge d'autres accélérateurs.

L'autograd et le graphe dynamique

L'entraînement exige le gradient de la perte par rapport à chaque paramètre. PyTorch l'obtient en enregistrant, pendant le calcul avant, les opérations effectuées sur les tenseurs : ce graphe est construit à la volée, à chaque exécution. L'appel à backward() le parcourt à l'envers. Comme le graphe suit simplement l'exécution du code Python, un modèle peut contenir des boucles ou des conditions, et une erreur se repère avec un débogueur classique.

La compilation

L'exécution impérative a un coût : chaque opération est lancée séparément. Depuis PyTorch 2.0, torch.compile capture le code d'un modèle et génère une version optimisée, tout en gardant l'écriture habituelle. C'est la réponse de PyTorch à l'approche de compilation de JAX et de TensorFlow.

L'exemple ci-dessous entraîne un petit réseau à reconnaître des chiffres manuscrits MNIST, avec torchvision pour les données. Le code utilise le GPU s'il est disponible.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Choisir le périphérique de calcul
device = "cuda" if torch.cuda.is_available() else "cpu"

# Charger MNIST et le découper en lots
data = datasets.MNIST(
    "data", train=True, download=True, transform=transforms.ToTensor()
)
loader = DataLoader(data, batch_size=64, shuffle=True)

# Définir le modèle comme une suite de couches
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28 * 28, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
).to(device)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(3):
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)

        loss = loss_fn(model(images), labels)  # calcul avant et perte
        optimizer.zero_grad()                  # remise à zéro des gradients
        loss.backward()                        # rétropropagation (autograd)
        optimizer.step()                       # mise à jour des poids

    print(f"Époque {epoch + 1} : perte = {loss.item():.4f}")

torch.save(model.state_dict(), "mnist.pt")

Le modèle et chaque lot sont envoyés sur le périphérique choisi avec .to(device). Les quatre lignes au cœur de la boucle suivent le cycle décrit plus haut : perte, remise à zéro des gradients, rétropropagation, mise à jour. Cette boucle explicite est la signature de PyTorch : elle demande quelques lignes de plus qu'une API de haut niveau, mais chaque étape reste visible et modifiable.

Les deux bibliothèques sont des frameworks complets de deep learning. TensorFlow, publié par Google en 2015, a dominé les premières années. PyTorch l'a progressivement rattrapé puis dépassé dans la recherche, avant de s'imposer aussi dans l'écosystème des LLM.

Ce n'est pas un choix binaire. Les concepts sont les mêmes : tenseurs, gradients, optimiseurs. Une équipe qui maîtrise l'un apprend l'autre en quelques jours. Le critère pratique est l'écosystème : si les modèles et le code à reprendre sont en PyTorch, c'est PyTorch.

JAX, développé par Google, constitue une troisième voie. Il applique des transformations (dérivation, compilation XLA, vectorisation) à des fonctions pures, et séduit pour la recherche à grande échelle et le calcul scientifique. Son écosystème reste toutefois plus restreint que celui de PyTorch.

Entraîner un modèle n'est qu'une partie du travail. Pour l'exploiter, quelques règles s'imposent.

• Fixer les versions : PyTorch, CUDA et les pilotes GPU doivent rester compatibles. Les images Docker officielles évitent la plupart des écarts.
• Exploiter le GPU : précision mixte, lots de taille adaptée et chargement de données en parallèle (num_workers) sont les premiers leviers de performance.
• Distribuer l'entraînement : DistributedDataParallel répartit un entraînement sur plusieurs GPU ou machines.
• Sauvegarder proprement : enregistrer le state_dict du modèle plutôt que l'objet complet, avec la version de la bibliothèque.
• Servir le modèle : choisir un moteur adapté, comme vLLM pour les LLM, ou une plateforme comme KServe sur Kubernetes.
• Suivre les expériences : un outil comme MLflow garde la trace des paramètres, des métriques et des modèles.