Skip to main content

​
Overview

The ONNX export module converts trained scikit-learn pipelines to the Open Neural Network Exchange (ONNX) format, enabling model deployment across different runtime environments and frameworks.

​
Function Signature

def export_pipeline_to_onnx(
    pipeline, 
    output_path: Path, 
    n_features: int
) -> bool:
    """
    Export a scikit-learn pipeline to ONNX format.
    
    Args:
        pipeline: Trained scikit-learn pipeline or model
        output_path: Path where ONNX model will be saved
        n_features: Number of input features
    
    Returns:
        True if export succeeded, False if an error occurred
    """
Source: deployment/onnx_export.py:6

​
Usage Example

from pathlib import Path
from deployment.onnx_export import export_pipeline_to_onnx
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier

# Create and train pipeline
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('classifier', RandomForestClassifier(n_estimators=100))
])
pipeline.fit(X_train, y_train)

# Export to ONNX
output_path = Path('models/hospital_risk_model.onnx')
success = export_pipeline_to_onnx(
    pipeline=pipeline,
    output_path=output_path,
    n_features=X_train.shape[1]
)

if success:
    print(f"Model exported to {output_path}")
else:
    print("Export failed - check dependencies")

​
Why ONNX?

ONNX provides several deployment advantages:

​
Cross-Platform Compatibility

  • Deploy scikit-learn models on platforms without Python runtime
  • Run models in C++, Java, C#, JavaScript environments
  • Execute on edge devices with ONNX Runtime

​
Performance Optimization

  • Hardware-specific optimizations (CPU, GPU, NPU)
  • Graph-level optimizations (operator fusion, constant folding)
  • Quantization support for reduced memory footprint

​
Interoperability

  • Convert between different ML frameworks
  • Standardized model format for MLOps pipelines
  • Version control and model lineage tracking

​
Implementation Details

The function performs the following steps:
  1. Import Dependencies: Uses skl2onnx library for conversion
  2. Define Input Schema: Creates FloatTensorType with dynamic batch size (None)
  3. Convert Pipeline: Transforms scikit-learn pipeline to ONNX graph
  4. Serialize: Writes protobuf-serialized model to disk
  5. Error Handling: Returns False on any exception
Source code from deployment/onnx_export.py:6-17: 
def export_pipeline_to_onnx(pipeline, output_path: Path, n_features: int) -> bool:
    try:
        from skl2onnx import convert_sklearn
        from skl2onnx.common.data_types import FloatTensorType

        initial_type = [("float_input", FloatTensorType([None, n_features]))]
        onx = convert_sklearn(pipeline, initial_types=initial_type)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_bytes(onx.SerializeToString())
        return True
    except Exception:
        return False

​
Dependencies

Required packages: 
pip install skl2onnx onnxruntime
Note: skl2onnx supports most common scikit-learn estimators. Check compatibility matrix for your specific pipeline components.

​
Input Type Definition

The input tensor is defined as: 
FloatTensorType([None, n_features])
  • First dimension (None): Dynamic batch size - allows variable-sized inputs
  • Second dimension (n_features): Fixed feature count matching training data
  • Data type: 32-bit floating point (fp32)

​
Inference with ONNX Runtime

After export, use ONNX Runtime for inference: 
import onnxruntime as rt
import numpy as np

# Load ONNX model
sess = rt.InferenceSession('models/hospital_risk_model.onnx')

# Prepare input (must match n_features)
X_input = np.array([[45, 120, 72]], dtype=np.float32)

# Run inference
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[1].name  # Index 1 for probabilities
probs = sess.run([output_name], {input_name: X_input})[0]

print(f"Risk probability: {probs[0][1]:.3f}")

​
Error Handling

The function returns False on any exception, including:
  • Missing dependencies: skl2onnx not installed
  • Unsupported operators: Pipeline contains non-convertible components
  • File system errors: Cannot write to output_path
  • Memory errors: Model too large for available RAM
For debugging, you can modify the function to log exceptions: 
import logging

try:
    # ... conversion code ...
    return True
except Exception as e:
    logging.error(f"ONNX export failed: {e}")
    return False

​
Model Validation

After export, validate ONNX model produces identical outputs: 
import numpy as np
import onnxruntime as rt

# Get scikit-learn predictions
sklearn_probs = pipeline.predict_proba(X_test)[:, 1]

# Get ONNX predictions
sess = rt.InferenceSession(str(output_path))
input_name = sess.get_inputs()[0].name
onnx_probs = sess.run(None, {input_name: X_test.values.astype(np.float32)})[1][:, 1]

# Compare
np.testing.assert_allclose(sklearn_probs, onnx_probs, rtol=1e-5)
print("✓ ONNX export validated")

​
Directory Creation

The function automatically creates parent directories: 
output_path.parent.mkdir(parents=True, exist_ok=True)
This ensures models/production/v2/model.onnx works even if intermediate directories don’t exist.

Build docs developers (and LLMs) love

Get started for freeTalk to us