How to export a scikit-learn model to ONNX

Model artifacts often need to leave the Python training stack before they can be served by smaller runtimes or non-Python applications. Exporting a supported scikit-learn estimator to ONNX creates a binary model file that ONNX Runtime can load without reconstructing the original Python estimator object.

skl2onnx handles the conversion layer for core scikit-learn estimators and pipelines that have registered converters. The converter needs a trained model plus a sample input or explicit input type information, because the ONNX graph must know the input name, data type, and feature shape before serving.

Use an isolated Python environment with the same feature order and numeric dtype used at training time. The smoke script trains a small LogisticRegression classifier, writes iris-logreg.onnx, loads the file through ONNX Runtime, and checks that the first labels match the original estimator.

Steps to export a scikit-learn model to ONNX:

  1. Install the converter and runtime packages in the active Python environment.
    $ python -m pip install --upgrade scikit-learn skl2onnx onnxruntime

    skl2onnx performs the conversion, and onnxruntime provides the CPU inference session used for the smoke test.

  2. Create export_iris_onnx.py in the project directory.
    export_iris_onnx.py
    from pathlib import Path
     
    import numpy as np
    import onnxruntime as ort
    from sklearn.datasets import load_iris
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import train_test_split
    from skl2onnx import to_onnx
     
    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, _ = train_test_split(
        X,
        y,
        test_size=0.2,
        random_state=42,
        stratify=y,
    )
     
    model = LogisticRegression(max_iter=500)
    model.fit(X_train, y_train)
     
    onnx_model = to_onnx(model, X_train[:1].astype(np.float32))
    onnx_path = Path("iris-logreg.onnx")
    onnx_path.write_bytes(onnx_model.SerializeToString())
     
    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    input_name = session.get_inputs()[0].name
    output_names = [output.name for output in session.get_outputs()]
    onnx_outputs = session.run(None, {input_name: X_test[:3].astype(np.float32)})
     
    sklearn_labels = model.predict(X_test[:3])
    onnx_labels = onnx_outputs[0]
     
    print(f"artifact: {onnx_path}")
    print(f"input: {input_name}")
    print(f"outputs: {output_names}")
    print(f"sklearn_labels: {sklearn_labels.tolist()}")
    print(f"onnx_labels: {onnx_labels.tolist()}")
    print(f"labels_match: {bool(np.array_equal(sklearn_labels, onnx_labels))}")

    The sample input passed to to_onnx() sets the graph input name, tensor type, and feature count. Keep the same column order that the trained estimator expects.

    skl2onnx does not convert every estimator. If conversion fails with a missing converter error, switch to a supported estimator or register a converter for the custom component before exporting.

  3. Run the export script.
    $ python export_iris_onnx.py
    artifact: iris-logreg.onnx
    input: X
    outputs: ['output_label', 'output_probability']
    sklearn_labels: [0, 2, 1]
    onnx_labels: [0, 2, 1]
    labels_match: True

    labels_match: True confirms that ONNX Runtime loaded the saved artifact and returned the same class labels for the smoke-test rows.

  4. Remove the sample script after moving the export code into the training project.
    $ rm export_iris_onnx.py