sklearn_logistic_regression_to_onnx_conversion_and_inference.py

1# Train a model.
2from sklearn.datasets import load_iris
3from sklearn.model_selection import train_test_split
4from sklearn.linear_model import LogisticRegression
5iris = load_iris()
6X, y = iris.data, iris.target
7X_train, X_test, y_train, y_test = train_test_split(X, y)
8clr = LogisticRegression()
9clr.fit(X_train, y_train)
10
11# Convert into ONNX format
12from skl2onnx import convert_sklearn
13from importlib_metadata import version
14from skl2onnx.common.data_types import FloatTensorType
15
16initial_type = [('float_input', FloatTensorType([None, 4]))]
17onx = convert_sklearn(clr, initial_types=initial_type)
18with open("logreg_iris.onnx", "wb") as f:
19    f.write(onx.SerializeToString())
20
21# Compute the prediction with onnxruntime
22import onnxruntime as rt
23import numpy
24sess = rt.InferenceSession("logreg_iris.onnx", providers=["CPUExecutionProvider"])
25input_name = sess.get_inputs()[0].name
26label_name = sess.get_outputs()[0].name
27pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0]