Aller au contenu principal

Module 2 — Export depuis PyTorch

Le cours PyTorch (module 10) présente torch.onnx.export à titre d'introduction. On va maintenant l'utiliser sérieusement, sur le fil rouge du cours : le ResNet18 entraîné sur Fashion-MNIST au module 9 de PyTorch. Ce module cadre les paramètres qui comptent, montre les pièges du traçage et introduit le nouvel exportateur Dynamo, qui devient le chemin par défaut depuis PyTorch 2.5.

Rappel : deux exportateurs coexistent​

Depuis PyTorch 2.1, deux mécanismes d'export cohabitent sous la même API.

  • L'exportateur historique, basé sur TorchScript, appelé simplement torch.onnx.export(...). Il est stable, très largement utilisé, bien documenté, mais construit son graphe par traçage — il exécute le modèle sur une entrée exemple et enregistre les opérations vues.
  • L'exportateur Dynamo, activé par dynamo=True. Il repose sur torch.compile et torch.export, capture directement le programme Python en un graphe symbolique et gère mieux le contrôle de flux, les formes dynamiques et les modèles récents.

Dans ce cours, on utilise l'exportateur historique par défaut — c'est celui qui existe partout, y compris dans les tutoriels des équipes NVIDIA et Microsoft — et l'on montre Dynamo à la fin comme trajectoire à suivre pour les nouveaux projets.

Un export minimal, ligne par ligne​

import torch
from torchvision.models import resnet18

# 1. Recharger le modèle entraîné du cours PyTorch
modele = resnet18(weights=None)
modele.fc = torch.nn.Linear(512, 10)
modele.load_state_dict(torch.load("resnet18_fashion.pt", map_location="cpu"))
modele.eval()

# 2. Fabriquer une entrée exemple avec la bonne forme et le bon dtype
exemple = torch.randn(1, 3, 224, 224)

# 3. Exporter
torch.onnx.export(
    modele,
    exemple,
    "resnet18_fashion.onnx",
    input_names=["entree"],
    output_names=["logits"],
    dynamic_axes={
        "entree": {0: "batch"},
        "logits": {0: "batch"},
    },
    opset_version=17,
    do_constant_folding=True,
)
print("Export terminé.")

Sept paramètres à comprendre.

modele doit être en mode eval(). Les couches Dropout ne font alors plus rien, BatchNorm utilise ses moyennes glissantes plutôt que la statistique du lot. Oublier eval() produit un modèle exporté qui prédit correctement puis dérive après plusieurs appels — le pire des cas, car il passe la vérification numérique sur un seul lot.

exemple est une entrée réaliste. Sa forme et son type conditionnent le traçage : un float32 de forme (1, 3, 224, 224) déclare un tenseur d'images RGB, un long de forme (1, 128) déclarerait des tokens. Les valeurs n'importent pas, mais l'entrée doit être exécutable par le modèle sans lever d'exception.

input_names et output_names donnent des noms parlants aux tenseurs d'entrée et de sortie. Sans eux, le graphe contient input.1, input.2, output.7, ce qui rend le service et le débogage pénibles. C'est un choix qu'on ne peut pas changer après coup facilement.

dynamic_axes est le paramètre le plus important après le modèle. Il déclare quels axes des tenseurs nommés sont variables à l'inférence. Sans lui, la forme de l'exemple est figée dans le graphe : un modèle exporté avec (1, 3, 224, 224) refuse un lot de 32 images. Pour Fashion-MNIST, la taille de lot est variable ; pour un modèle NLP, la longueur de séquence l'est aussi.

opset_version=17 fixe le jeu d'opérateurs. Le module 1 explique pourquoi c'est important ; en pratique, on choisit la version la plus ancienne qui supporte les opérateurs du modèle et est reconnue par le runtime cible. TensorRT 10 accepte jusqu'à l'opset 21, ONNX Runtime 1.19 supporte tout jusqu'à l'opset 22, donc 17 laisse une marge confortable.

do_constant_folding=True demande à PyTorch de précalculer les sous-expressions dont les entrées sont toutes constantes. Une convolution avec un noyau et un biais constants est repliée en une seule constante. C'est un optimisation gratuite qu'on active toujours.

Les axes dynamiques en détail​

dynamic_axes accepte un dictionnaire dont les clés sont les noms de tenseurs et les valeurs des dictionnaires {index_axe: nom_symbolique}. Un modèle qui accepte à la fois une taille de lot et une hauteur variable s'écrit ainsi :

dynamic_axes={
    "entree": {0: "batch", 2: "hauteur", 3: "largeur"},
    "logits": {0: "batch"},
}

Le nom symbolique n'a pas d'importance intrinsèque — il apparaît dans les métadonnées du graphe — mais deux axes portant le même nom doivent rester égaux à l'inférence. Si entree et masque déclarent tous deux "longueur" pour leur axe 1, tenter d'appeler la session avec deux tenseurs de longueurs différentes lèvera une erreur explicite. C'est un contrat entre l'exportateur et le runtime.

Axes dynamiques et graphes figés

Un axe non déclaré est figé. Un export sans dynamic_axes produit un modèle qui n'accepte que la forme exacte de l'exemple. C'est la première cause d'échec en production : le service qui essaie d'envoyer un lot de 8 images sur un modèle exporté avec (1, 3, 224, 224) reçoit un InvalidArgument. Toujours déclarer dynamic_axes, même pour un modèle qu'on pense servir un exemple à la fois : les mesures de latence du module 8 exigent des lots.

Le piège du traçage sur le contrôle de flux​

L'exportateur historique trace le modèle sur l'entrée exemple. Les branchements if ne sont pas capturés : seule la branche prise sur cet exemple précis apparaît dans le graphe. Un modèle qui applique une normalisation différente selon x.max() produit un ONNX qui applique toujours celle observée pendant le traçage.

class ModeleFragile(torch.nn.Module):
    def forward(self, x):
        if x.max() > 0.5:                # branchement fragile au traçage
            return x * 2
        return x + 1

L'exportateur historique lève un TracerWarning sur ce type de code — le lire est obligatoire. En cas de doute, deux réponses : soit réécrire le modèle sans contrôle de flux dépendant des données (via torch.where), soit passer à Dynamo qui gère le contrôle de flux nativement.

L'exportateur Dynamo​

Depuis PyTorch 2.5, on peut activer le nouvel exportateur en une ligne :

export_program = torch.onnx.export(
    modele,
    (exemple,),
    "resnet18_fashion_dynamo.onnx",
    dynamo=True,
    dynamic_shapes={
        "entree": {0: torch.export.Dim("batch")},
    },
    opset_version=18,
)

Trois différences visibles. Le second argument devient un tuple de tenseurs. dynamic_shapes remplace dynamic_axes avec une API plus explicite : les axes symboliques sont des objets torch.export.Dim. Enfin, Dynamo capture le programme via torch.export.export, donc le contrôle de flux structurel est préservé, et l'exportateur produit souvent un graphe plus propre.

Dynamo est encore jeune. Un modèle complexe qui utilise des opérations peu courantes peut aujourd'hui casser à l'export ; l'exportateur historique reste le repli. La consigne pratique pour 2026 : essayer Dynamo d'abord ; si l'export échoue avec une trace opaque, revenir à torch.onnx.export classique.

Relire l'export avec onnx​

Après tout export, on relit le fichier pour vérifier ce qu'on a produit.

import onnx

modele_onnx = onnx.load("resnet18_fashion.onnx")
onnx.checker.check_model(modele_onnx)

print("Opset :", modele_onnx.opset_import[0].version)
for e in modele_onnx.graph.input:
    axes = [d.dim_value or d.dim_param for d in e.type.tensor_type.shape.dim]
    print("Entrée", e.name, ":", axes)

Si l'opset ne correspond pas à celui demandé, ou si un axe apparaît figé alors qu'on l'a déclaré dynamique, le problème est dans l'export lui-même, pas dans le runtime. On corrige avant d'aller plus loin.

Le cas du petit encodeur de texte​

Le second modèle du fil rouge — un petit encodeur qui classe des phrases en cinq catégories — s'exporte de manière similaire, avec deux entrées : les identifiants de tokens et le masque d'attention.

tokens = torch.randint(0, 30000, (1, 128))
masque = torch.ones(1, 128, dtype=torch.long)

torch.onnx.export(
    encodeur_texte,
    (tokens, masque),
    "encodeur_texte.onnx",
    input_names=["tokens", "masque"],
    output_names=["logits"],
    dynamic_axes={
        "tokens":  {0: "batch", 1: "longueur"},
        "masque":  {0: "batch", 1: "longueur"},
        "logits":  {0: "batch"},
    },
    opset_version=17,
)

Ce modèle sera repris au module 9 : son opérateur d'attention n'est pas encore stable en opset 17 sur toutes les cibles, et il faudra choisir entre monter en opset 20 ou décomposer manuellement l'attention.

En résumé​

  • Un export PyTorch réussit en mode eval() avec une entrée exemple bien typée ; oublier eval() produit une dérive silencieuse due à BatchNorm ou Dropout mal figés.
  • dynamic_axes est obligatoire dès qu'on veut varier la taille de lot ou la longueur de séquence à l'inférence ; un axe non déclaré est figé, sans avertissement.
  • L'exportateur historique trace le modèle sur une entrée : le contrôle de flux dépendant des données n'est pas capturé, et il faut lire les TracerWarning.
  • L'exportateur Dynamo (dynamo=True) capture le programme et gère mieux le contrôle de flux ; c'est la trajectoire moderne mais elle reste plus fragile sur les modèles atypiques.

Le module suivant fait la même opération depuis TensorFlow et Keras : on y verra une convention d'axes différente (NHWC contre NCHW) et une chaîne d'outils distincte (tf2onnx).