How to use ONNX in ML Apps

Photo by Ries Bosch on Unsplash

To export a Logistic Regression model from scikit-learn to ONNX format, you need to use the sklearn-onnx library, which provides a straightforward way to convert scikit-learn models into ONNX format. First, ensure you have the necessary packages installed:

pip install skl2onnx onnxruntime numpy scikit-learn

Here’s how you can train a simple Logistic Regression model with scikit-learn and export it to ONNX format:

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
import numpy as np
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# Load example data
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Train a logistic regression model
model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

# Convert to ONNX format
# The initial type is necessary for the conversion, it specifies the type and shape of the input
# Here, 'None' indicates that the first dimension is variable, and 4 is the number of features in the iris dataset
initial_type = [('float_input', FloatTensorType([None, 4]))]
onnx_model = convert_sklearn(model, initial_types=initial_type)

# Save the model to a file
with open("logistic_regression.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

print("Model successfully exported to logistic_regression.onnx")

This code performs the following steps:

1. Loads the Iris dataset and splits it into training and testing sets.
2. Trains a Logistic Regression model on the training data.
3. Converts the trained model to ONNX format, specifying the input’s data type and shape (initial_type). This is crucial because ONNX needs to know the size and type of the input tensor.
4. Saves the ONNX model to a file named logistic_regression.onnx.

Keep in mind that skl2onnx supports a variety of scikit-learn models, but there could be limitations depending on the specific version of skl2onnx and the complexity of the scikit-learn model you are trying to convert.

Load the model:

To load an ONNX model and make predictions with it, you can use the onnxruntime library, which provides an efficient and easy way to run inference on ONNX models. Here's how you can do it, continuing from the example where we exported a Logistic Regression model to ONNX format:

import onnxruntime as ort
import numpy as np

# Load the ONNX model using ONNX Runtime
sess = ort.InferenceSession("logistic_regression.onnx")

# Prepare a sample input for prediction
# For example, using the first sample from the X_test set previously created
sample_input = X_test[0].reshape(1, -1).astype(np.float32)  # Reshape and ensure type is float32

# Prepare the input for the model (name must match the input name from the export step)
input_name = sess.get_inputs()[0].name
input_data = {input_name: sample_input}

# Run the model (inference)
predictions = sess.run(None, input_data)

# The output of the model
print("Predicted class:", np.argmax(predictions[1]), "with probabilities:", predictions[1])

# If you want to see the raw output
print("Raw model output:", predictions)

This code snippet does the following:

1. Initializes an inference session using onnxruntime with the exported ONNX model.
2. Prepares a sample input for prediction. Here, we’re using a sample from X_test, reshaping it to match the expected input shape for the model, and ensuring it's of type float32, as ONNX models are particular about input types and shapes.
3. Retrieves the input name from the model session to correctly map the input data for the prediction. This is important because ONNX models expect a dictionary with input names as keys.
4. Runs the model with the input data using the sess.run method. This method returns a list of outputs, and since Logistic Regression models output probabilities for each class, we can use these to determine the predicted class by finding the index of the highest probability.
5. Prints the predicted class by taking the argmax of the probabilities. The output also includes the raw probabilities for each class to see the model's confidence levels.

This example assumes you have a trained Logistic Regression model saved as “logistic_regression.onnx” and that you’ve followed the steps to export the model as described in the previous example. The actual output will depend on the input sample and the trained model.