Aller au contenu principal

Module 3 — Export depuis TensorFlow

Le module 2 a couvert la chaîne PyTorch. TensorFlow suit une logique différente, et une équipe qui hérite d'un modèle Keras entraîné il y a deux ans a rarement la latitude de le réentraîner sous PyTorch juste pour l'exporter en ONNX. Ce module montre comment convertir sans réentraîner, avec l'outil de référence : tf2onnx.

Le paysage TensorFlow​

Un modèle TensorFlow peut vivre sous plusieurs formes.

  • Keras .h5 ou .keras : le format classique de l'API haute, contient l'architecture et les poids.
  • SavedModel : le format universel de TensorFlow, un dossier qui contient un graphe protobuf et un dossier variables/ pour les poids. C'est le format que Google recommande depuis TF 2.
  • Frozen graph .pb : ancien, monolithique, encore présent dans des dépôts hérités.
  • TFLite .tflite : orienté mobile ; hors du périmètre de ce module.

tf2onnx accepte le SavedModel et le Keras. Le frozen graph reste possible avec une commande supplémentaire, mais c'est la voie la moins fiable et il vaut mieux le reconvertir en SavedModel d'abord.

Installer et lancer tf2onnx​

tf2onnx est un paquet Python distinct de TensorFlow, avec un outil en ligne de commande et une API Python. La version 1.16 et plus supporte les opsets jusqu'à 19 et TensorFlow 2.15.

pip install tf2onnx onnx tensorflow==2.15

Le chemin le plus court, pour un SavedModel déjà sur disque, est une commande shell :

python -m tf2onnx.convert \
  --saved-model ./modele_savedmodel \
  --output modele.onnx \
  --opset 17

Pour un modèle Keras :

python -m tf2onnx.convert \
  --keras modele.keras \
  --output modele.onnx \
  --opset 17

L'outil charge le modèle, retrace le graphe, mappe chaque opérateur TensorFlow à son équivalent ONNX quand il existe, écrit le fichier. Un rapport en fin d'exécution liste les opérateurs convertis et signale ceux qui n'ont pas d'équivalent.

L'API Python : plus de contrôle​

La ligne de commande suffit pour les cas courants. Dès que l'on veut renommer des tenseurs, déclarer des axes dynamiques ou exporter à partir d'un modèle Keras chargé en mémoire, l'API Python est plus lisible.

import tensorflow as tf
import tf2onnx
import onnx

# Un modèle Keras factice qui joue le rôle de « ResNet18 côté TensorFlow »
modele_keras = tf.keras.applications.ResNet50(
    weights="imagenet",
    input_shape=(224, 224, 3),
    include_top=True,
)

# Déclarer la signature d'appel : forme, type, nom
signature = [tf.TensorSpec([None, 224, 224, 3], tf.float32, name="entree")]

modele_onnx, _ = tf2onnx.convert.from_keras(
    modele_keras,
    input_signature=signature,
    opset=17,
    output_path="resnet50_tf.onnx",
)

onnx.checker.check_model(modele_onnx)
print("Entrées :", [e.name for e in modele_onnx.graph.input])
print("Sorties :", [s.name for s in modele_onnx.graph.output])

Trois points à noter. La TensorSpec déclare une dimension None en tête pour la taille de lot — c'est l'équivalent TensorFlow des axes dynamiques ONNX. L'argument opset choisit la version de jeu d'opérateurs, comme dans PyTorch. Le retour de from_keras est un couple (modele_onnx, external_tensor_storage) où le second sert pour le stockage externe des poids d'un très gros modèle.

Signatures d'un SavedModel​

Un SavedModel peut exposer plusieurs signatures — plusieurs points d'entrée nommés. C'est courant pour un modèle de recommandation qui expose à la fois predict et train, ou pour un transformeur qui expose encode et generate. tf2onnx exporte une signature à la fois ; le paramètre --signature_def (ou signature_def en API) la sélectionne.

python -m tf2onnx.convert \
  --saved-model ./modele_savedmodel \
  --signature_def serving_default \
  --output modele.onnx \
  --opset 17

Sans ce paramètre, tf2onnx prend serving_default si elle existe, sinon lève une erreur. Un SavedModel qu'on ne connaît pas se lit d'abord avec saved_model_cli show --dir ./modele_savedmodel --all — la commande listera les signatures et leurs entrées et sorties.

NHWC contre NCHW : la conversion silencieuse​

TensorFlow suit par défaut la convention NHWC — (lot, hauteur, largeur, canaux) — héritée de son passé sur CPU. ONNX, comme PyTorch et CUDA, préfère NCHW — (lot, canaux, hauteur, largeur).

tf2onnx convertit automatiquement. Le graphe résultant est en NCHW, et l'outil insère des opérateurs Transpose autour des convolutions pour respecter cette convention en interne. Sur ONNX Runtime CPU, cette conversion peut aussi coûter un peu, car les kernels préfèrent parfois NHWC.

Deux implications pratiques.

  • Le tenseur d'entrée attendu par le modèle ONNX exporté est en NCHW par défaut : (N, 3, 224, 224) et non (N, 224, 224, 3). Un service qui envoie du NHWC produit du bruit sans erreur explicite, parce que les valeurs restent des float32 valides.
  • Si vous préférez rester en NHWC — utile pour un service qui reçoit directement des images non permutées — l'option --inputs-as-nchw entree demande à tf2onnx de générer un modèle NHWC natif, en s'appuyant sur les opérateurs ONNX qui l'acceptent (Conv, MaxPool, avec l'attribut channels_last).
python -m tf2onnx.convert \
  --saved-model ./modele_savedmodel \
  --output modele.onnx \
  --opset 17 \
  --inputs-as-nchw entree
La convention se propage à toutes les mesures

Une équipe qui a mesuré la latence de son modèle TensorFlow en NHWC et compare à son export ONNX en NCHW compare deux tailles de tenseurs différentes, et attribue au format un écart qui vient en réalité des permutations. Standardiser la convention d'un bout à l'autre du protocole de mesure est la seule façon d'obtenir un chiffre reproductible (module 8).

Les opérateurs qui ne passent pas​

tf2onnx couvre l'écrasante majorité des opérateurs TensorFlow. Il en reste qui n'ont pas d'équivalent ONNX direct ; l'outil les signale à la fin de la conversion sous la forme d'un rapport :

ERROR: Tensorflow op [MyCustomOp: MyCustomOp] is not supported

Trois réponses classiques.

  • Décomposer l'opérateur en fonctions primitives, quand c'est possible : un RandomShuffle peut être réécrit avec tf.gather et un permutateur figé.
  • Monter en opset : un opérateur qui n'existe pas en opset 13 peut apparaître en opset 17.
  • Écrire une fonction ONNX personnalisée en dernier recours (module 9).

Un modèle avec un préprocesseur personnalisé — décodage JPEG, extraction de spectrogrammes — n'est en général pas convertible tel quel. La règle : exporter en ONNX le modèle numérique nu, et laisser le préprocesseur au service (module 10).

Le cas Keras pur : signatures implicites​

Un modèle Keras qui n'a jamais été sauvegardé en SavedModel n'a pas de signature explicite. from_keras en fabrique une à partir de la première couche d'entrée. Si le modèle possède plusieurs entrées — par exemple un modèle multimodal texte + image — il faut lui passer une liste de TensorSpec dans l'ordre attendu par le call.

signature = [
    tf.TensorSpec([None, 224, 224, 3], tf.float32, name="image"),
    tf.TensorSpec([None, 128],         tf.int32,   name="tokens"),
]
modele_onnx, _ = tf2onnx.convert.from_keras(
    modele_multimodal,
    input_signature=signature,
    opset=17,
)

Sans input_signature, tf2onnx voit x[0] et x[1] avec des noms indéfinis. C'est convertible mais peu maintenable ; on préfère toujours les noms explicites.

Comparer un export TensorFlow et son homologue PyTorch​

Si les deux modèles du fil rouge existent dans les deux cadres — c'est le cas du ResNet50 dont les poids sont partagés — comparer les deux exports ONNX sur la même entrée révèle des dérives d'ordre 1e-5 en float32, essentiellement dues aux implémentations différentes de BatchNorm. C'est le sujet du module 4.

En résumé​

  • tf2onnx convertit SavedModel et Keras vers ONNX ; la ligne de commande suffit pour un cas simple, l'API Python devient nécessaire pour renommer les tenseurs et déclarer des axes dynamiques via TensorSpec.
  • Un SavedModel expose souvent plusieurs signatures ; en sélectionner explicitement une avec --signature_def évite les erreurs muettes de conversion.
  • tf2onnx convertit NHWC en NCHW par défaut, en insérant des Transpose autour des convolutions ; garder NHWC natif est possible avec --inputs-as-nchw, mais il faut alors être cohérent d'un bout à l'autre.
  • Un opérateur non pris en charge s'affiche dans le rapport de fin ; les issues classiques sont la décomposition, la montée d'opset ou une fonction personnalisée (module 9).

Le module suivant part de ces deux exports — le ResNet PyTorch et le ResNet50 TensorFlow — et vérifie qu'ils rendent numériquement ce que rendaient les modèles d'origine.