Note
Click here to download the full example code
Converts a scikit-learn model¶
The gallery of sklearn-onnx provides many examples with scikit-learn. The following takes the same example and rewrites it with onnxmltools.
Train and deploy a model usually involves the three following steps:
train a pipeline with scikit-learn,
convert it into ONNX with sklearn-onnx,
predict with onnxruntime.
Train a model¶
A very basic example using random forest and the iris dataset.
import numpy
import onnx
import sklearn
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
import onnxruntime as rt
import skl2onnx
import onnxmltools
from onnxconverter_common.data_types import FloatTensorType
from onnxmltools.convert import convert_sklearn
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
clr = RandomForestClassifier()
clr.fit(X_train, y_train)
print(clr)
Out:
RandomForestClassifier()
Convert a model into ONNX¶
initial_type = [('float_input', FloatTensorType([None, 4]))]
onx = convert_sklearn(clr, initial_types=initial_type)
with open("rf_iris.onnx", "wb") as f:
f.write(onx.SerializeToString())
Compute the prediction with onnxruntime¶
sess = rt.InferenceSession("rf_iris.onnx")
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run(
[label_name], {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)
Out:
[1 1 0 1 0 2 2 0 1 1 2 1 2 2 1 1 1 2 2 2 2 0 2 2 0 0 1 1 0 0 2 1 1 0 0 0 2
1]
Full example with a logistic regression
clr = LogisticRegression()
clr.fit(X_train, y_train)
initial_type = [('float_input', FloatTensorType([None, X_train.shape[1]]))]
onx = convert_sklearn(clr, initial_types=initial_type)
with open("logreg_iris.onnx", "wb") as f:
f.write(onx.SerializeToString())
sess = rt.InferenceSession("logreg_iris.onnx")
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name],
{input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)
Out:
somewhereonnxmltools-jenkins_39_std/_venv/lib/python3.9/site-packages/sklearn/linear_model/_logistic.py:814: ConvergenceWarning: lbfgs failed to converge (status=1):
STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.
Increase the number of iterations (max_iter) or scale the data as shown in:
https://scikit-learn.org/stable/modules/preprocessing.html
Please also refer to the documentation for alternative solver options:
https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression
n_iter_i = _check_optimize_result(
[1 2 0 1 0 2 2 0 1 1 2 1 2 2 1 1 1 2 2 2 2 0 2 2 0 0 2 1 0 0 2 1 1 0 0 0 2
1]
Display the ONNX graph¶
Finally, let’s see the graph converted with onnxmltools.
import os
import matplotlib.pyplot as plt
from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer
pydot_graph = GetPydotGraph(
onx.graph, name=onx.graph.name, rankdir="TB",
node_producer=GetOpNodeProducer(
"docstring", color="yellow", fillcolor="yellow", style="filled"))
pydot_graph.write_dot("model.dot")
os.system('dot -O -Gdpi=300 -Tpng model.dot')
image = plt.imread("model.dot.png")
fig, ax = plt.subplots(figsize=(40, 20))
ax.imshow(image)
ax.axis('off')
Out:
(-0.5, 1982.5, 2558.5, -0.5)
Versions used for this example
print("numpy:", numpy.__version__)
print("scikit-learn:", sklearn.__version__)
print("onnx: ", onnx.__version__)
print("onnxruntime: ", rt.__version__)
print("onnxmltools: ", onnxmltools.__version__)
print("skl2onnx: ", skl2onnx.__version__)
Out:
numpy: 1.21.4
scikit-learn: 1.1.dev0
onnx: 1.10.2
onnxruntime: 1.10.91
onnxmltools: 1.10.0
skl2onnx: 1.10.4
Total running time of the script: ( 0 minutes 6.246 seconds)