机器学习模型转ONNX例子

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

def build_model(X_train, y_train):
    clr = LogisticRegression()
    clr.fit(X_train, y_train)
    return clr

def convert_model(clr):
    initial_type = [('float_input', FloatTensorType([None, 4]))]
    onx = convert_sklearn(clr, initial_types=initial_type)
    with open("logreg_iris.onnx", "wb") as f:
        f.write(onx.SerializeToString())

def run_model():
    sess = rt.InferenceSession(
        "logreg_iris.onnx", providers=rt.get_available_providers()) # ['CUDAExecutionProvider', 'CPUExecutionProvider']
    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)

if __name__ == '__main__':
    iris = load_iris()
    X, y = iris.data, iris.target
    X_train, X_test, y_train, y_test = train_test_split(X, y)
    clr = build_model(X_train, y_train)
    convert_model(clr)
    run_model()