PyTorch : choisir eval, no_grad et inference_mode

Dans PyTorch, model.eval() change le comportement des couches du modèle. torch.no_grad() et torch.inference_mode() contrôlent le travail lié aux gradients. Pour une prédiction ordinaire, il faut généralement choisir les deux : le comportement d’évaluation du modèle et le mode de gradient adapté au passage avant.

Cette distinction compte lorsque la validation modifie les statistiques BatchNorm, que les prédictions conservent un graphe de calcul ou que des caractéristiques figées alimentent ensuite une couche entraînable. Les exemples ci-dessous ont été exécutés sur CPU avec PyTorch 2.15.0. Ils vérifient le comportement, sans mesurer la vitesse GPU ni les économies de mémoire.

Choisir la combinaison selon la tâche

Tâche	Comportement du modèle	Mode de gradient
Entraîner les paramètres	Généralement model.train()	Enregistrement normal des gradients
Prédire sans rétropropagation ultérieure	model.eval()	inference_mode() est un point de départ adapté
Valider sans gradients	model.eval()	no_grad(), ou inference_mode() si ses restrictions conviennent
Extraire des caractéristiques figées pour une tête entraînable	backbone.eval() ; head.train()	no_grad() autour du backbone seulement
Calculer un gradient par rapport à une entrée	Souvent model.eval()	Enregistrement normal ; l’entrée demande des gradients

Les notes autograd de PyTorch distinguent l’évaluation d’un module des modes de gradient. Ce sont deux contrôles distincts. Un modèle peut être en mode évaluation pendant qu’autograd enregistre ses opérations.

Ce que change model.eval()

Module.eval() équivaut à train(False). Il modifie l’indicateur d’entraînement du module et de ses sous-modules. Les couches qui utilisent cet indicateur peuvent alors changer de comportement. Il ne désactive pas les gradients et ne passe pas les indicateurs requires_grad des paramètres à false.

Dropout masque aléatoirement des valeurs pendant l’entraînement et devient une identité en évaluation. BatchNorm met normalement à jour ses statistiques cumulées pendant l’entraînement et les utilise en évaluation. Avec track_running_stats=False, il utilise les statistiques du lot dans les deux modes.

Nous avons vérifié ces comportements : en mode entraînement, Dropout masquait encore des valeurs dans no_grad() et BatchNorm modifiait encore sa moyenne cumulée. Dans les configurations par défaut, eval() changeait ces comportements. Désactiver les gradients ne suffit donc pas à préparer ces couches pour une validation ordinaire.

Un petit exemple de prédiction

Exécutez ce script Python autonome dans un environnement où PyTorch est installé. Les poids sont aléatoires et le modèle n’est pas entraîné. L’exemple montre l’état d’exécution, pas des prédictions utiles.

import torch
from torch import nn

torch.manual_seed(7)
model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(8, 2),
)
inputs = torch.randn(3, 4)

model.eval()
with torch.inference_mode():
    logits = model(inputs)

print(logits.shape, logits.requires_grad, torch.is_inference(logits))

La sortie observée était torch.Size([3, 2]) False True : trois lignes, aucune demande de gradient et un tenseur d’inférence. Dans une application réelle, chargez d’abord les poids entraînés et placez le modèle et les entrées sur des périphériques et types compatibles. Notre guide de sauvegarde et de chargement PyTorch traite les checkpoints.

inference_mode() désactive l’enregistrement des gradients et certaines opérations de suivi d’autograd, dont le suivi des vues et les compteurs de version. Il n’appelle pas eval(). Les tenseurs créés dans ce contexte ont aussi des restrictions d’utilisation ultérieure avec autograd.

Observer directement les différences

Ajoutez ce code au script de prédiction. Ses paramètres demandent toujours des gradients : la ligne avec eval seul enregistre donc un graphe, même si l’entrée ne demande pas de gradient.

from contextlib import nullcontext

cases = [
    ("train + no_grad", True, torch.no_grad),
    ("eval + default grad", False, nullcontext),
    ("eval + no_grad", False, torch.no_grad),
    ("eval + inference_mode", False, torch.inference_mode),
]

for label, training, context in cases:
    model.train(training)
    with context():
        output = model(inputs)
    print(label, model.training, output.requires_grad,
          torch.is_inference(output))

Cas	model.training	output.requires_grad	Tenseur d’inférence
train + no_grad	True	False	False
eval + default grad	False	True	False
eval + no_grad	False	False	False
eval + inference_mode	False	False	True

Le test CPU a reproduit les quatre lignes. Le modèle évalué a également produit des sorties correspondantes sous no_grad() et inference_mode(). Cela ne garantit pas des résultats identiques pour tous les modèles, périphériques ou opérations, et ne constitue pas une comparaison de performance.

Quand préférer no_grad()

no_grad() désactive l’enregistrement des gradients en mode inverse dans son contexte, tout en produisant des tenseurs ordinaires. Il convient lorsqu’un résultat devient ensuite l’entrée d’un calcul entraînable.

Cet exemple autonome fige un extracteur de caractéristiques, puis calcule les gradients d’une tête de classification distincte :

import torch
from torch import nn
from torch.nn import functional as F

torch.manual_seed(7)
backbone = nn.Sequential(nn.Linear(4, 8), nn.ReLU())
head = nn.Linear(8, 2)
backbone.requires_grad_(False)
backbone.eval()
head.train()

inputs = torch.randn(3, 4)
targets = torch.tensor([0, 1, 0])
with torch.no_grad():
    features = backbone(inputs)

loss = F.cross_entropy(head(features), targets)
loss.backward()

print(head.weight.grad is not None)
print(all(p.grad is None for p in backbone.parameters()))

Les deux vérifications ont affiché True. La tête a reçu des gradients, tandis que les paramètres neufs du backbone n’en ont pas reçu. Il s’agit d’une rétropropagation, sans mise à jour par un optimiseur. Une boucle réelle doit aussi effacer les gradients et effectuer l’étape d’optimisation ; notre guide d’accumulation des gradients présente ces responsabilités.

En remplaçant uniquement le contexte d’extraction par inference_mode(), la tête linéaire testée a rejeté le tenseur obtenu, car elle devait conserver cette entrée pour la rétropropagation. L’entrée n’a pas besoin de demander des gradients pour que la tête ait besoin de ses valeurs afin de calculer les gradients de ses poids.

Si des tenseurs d’inférence doivent franchir cette limite, un clonage hors du mode d’inférence peut produire des tenseurs ordinaires. Nous avons vérifié cette solution sur l’exemple. Elle ajoute une copie ; choisir no_grad() à la sortie de l’extracteur est souvent plus clair. Le clonage ne rétablit pas un chemin de gradient à travers des calculs qui n’ont jamais été enregistrés.

Conserver les gradients pour analyser une entrée

Certaines tâches d’évaluation nécessitent des dérivées, par exemple pour examiner comment un score varie avec son entrée. Le mode évaluation peut convenir, mais désactiver les gradients autour du passage avant supprimerait le chemin nécessaire.

import torch
from torch import nn

torch.manual_seed(7)
model = nn.Linear(4, 2)
model.requires_grad_(False)
model.eval()
inputs = torch.randn(1, 4, requires_grad=True)

score = model(inputs)[0, 0]
(input_gradient,) = torch.autograd.grad(score, inputs)
print(input_gradient.shape)

L’exemple a renvoyé torch.Size([1, 4]). Nous avons aussi vérifié que le gradient correspond à la première ligne de la matrice de poids. Les paramètres sont figés, mais une dérivée peut toujours être calculée par rapport à l’entrée. torch.autograd.grad renvoie directement cette dérivée.

Rétablir volontairement le comportement d’entraînement

Après la validation dans une boucle d’entraînement, appelez model.train() avant le passage avant suivant. Quitter un contexte no-grad ou inference rétablit l’état précédent du contexte ; cela n’annule pas l’appel à model.eval().

Avec un backbone figé et une tête entraînable, définissez explicitement les états des sous-modules. Un appel à train() sur le parent modifie aussi ses enfants et peut annuler le choix du mode évaluation pour le backbone. L’indicateur du module principal ne suffit pas à décrire tous les modèles à modes mixtes.

Les contextes de gradient sont locaux à chaque thread. Placez le contexte adapté autour de l’exécution dans chaque worker concerné. L’indicateur d’entraînement appartient au modèle : des requêtes de prédiction concurrentes ne doivent pas faire basculer un modèle partagé entre train et eval. Configurez son état prévu lors de l’initialisation.

Erreurs courantes à vérifier

• La validation modifie les statistiques : examinez les indicateurs d’entraînement, notamment BatchNorm, pas seulement le contexte de gradient.
• Les prédictions conservent des graphes : vérifiez l’enregistrement des gradients et les listes ou caches qui gardent ces sorties en mémoire.
• Une tête entraînable rejette les caractéristiques d’inférence : utilisez no_grad pour l’extraction ou créez une copie ordinaire hors du mode d’inférence.
• L’analyse du gradient d’entrée échoue : vérifiez l’indicateur de gradient de l’entrée et le contexte du passage avant.
• L’entraînement change après validation : rétablissez les états prévus du modèle et de ses sous-modules.

detach() sépare un tenseur de son graphe tout en partageant son stockage. Détacher une sortie après le passage avant n’empêche pas le travail de construction du graphe déjà effectué. Si la mémoire d’inférence continue d’augmenter, notre guide CUDA out of memory examine les sorties conservées et d’autres causes.

Deux détails concernent le code spécialisé : no_grad() comporte une exception pour les fonctions de création avec une option explicite de gradient et ne désactive pas la différentiation automatique en mode direct. Le mode d’inférence désactive aussi cette dernière. Mesurez la charge réelle avant d’annoncer un gain de vitesse ou de mémoire ; les petits tests CPU présentés ici vérifient uniquement le comportement.