Machine Learning

Exporter un pipeline complet en ONNX et vérifier la parité

Conversion du pipeline scikit-learn entier (preprocessing inclus) en ONNX, puis assertion numérique entre les sorties sklearn et onnxruntime — l'étape que l'on regrette d'avoir sautée.

Cas d'usage

Déploiement léger (sans sklearn) dans un service C#/C++/Rust ou un environnement contraint.

Prérequis

skl2onnx, onnxruntime, numpy

Python
import numpy as np
import onnxruntime as rt
from skl2onnx import to_onnx

X_sample = X_train.to_numpy().astype(np.float32)[:1]
onnx_model = to_onnx(
    pipe, X_sample,
    options={id(pipe.steps[-1][1]): {"zipmap": False}},  # sortie ndarray
)
with open("pipeline.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

sess = rt.InferenceSession("pipeline.onnx",
                           providers=["CPUExecutionProvider"])
entree = sess.get_inputs()[0].name
X_np = X_test.to_numpy().astype(np.float32)

proba_onnx = sess.run(None, {entree: X_np})[1]
proba_skl = pipe.predict_proba(X_test)

ecart = float(np.abs(proba_skl - proba_onnx).max())
assert ecart < 1e-4, f"Divergence sklearn/ONNX : {ecart}"
print(f"Parité OK (écart max {ecart:.2e})")

Résultat

Parité OK (écart max 3.42e-07)
>>> sess.get_inputs()[0].shape, sess.get_inputs()[0].type
([None, 24], 'tensor(float)')
>>> import os; os.path.getsize("pipeline.onnx") // 1024   # Ko sur disque
187
ONNXExportParitéProduction

Snippets liés

Retour au Data Lab