← retour aux snippets

ONNX: exporter un modèle scikit-learn

convertir un pipeline scikit-learn en ONNX pour l'inférence

python deploy #onnx#skl2onnx#export

ONNX: exporter un modèle scikit-learn

objectif

Expliquer et montrer comment convertir un pipeline scikit-learn en ONNX pour l’inférence.

code minimal

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

pipe = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))
X_dummy = [[0.0, 1.0, -1.0]]
pipe.fit([[0,0,0],[1,1,1],[2,2,2]],[0,1,1])

initial_types = [("input", FloatTensorType([None, 3]))]
onnx_model = convert_sklearn(pipe, initial_types=initial_types)
with open("model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

utilisation

# exécution avec onnxruntime
# import onnxruntime as rt
# sess = rt.InferenceSession("model.onnx")
# sess.run(None, {"input": np.array(X_dummy, dtype=np.float32)})

variante(s) utile(s)

# gérer les catégories via OneHotEncoder + OrdinalEncoder compatibles ONNX

notes

  • Assurez-vous que chaque étape du pipeline est supportée par skl2onnx.
  • Définissez correctement initial_types (n_features).