【sklearn下的SVM转onnx】

SVM模型转onnx

import joblib
from skl2onnx.common.data_types import *
import onnxmltools
 
# Update the input name and path for your sklearn model
input_skl_model = './model.pkl'
 
# input data type for your sklearn model
# 输入是Laplace算子，可以通过运行test_onnx.py报错来判断参数类型
input_data_type = [('int64_input', Int64TensorType([1, 2]))] 
 

# Change this path to the output name and path for the ONNX model
output_onnx_model = 'model.onnx'
 
# Load your sklearn model
skl_model = joblib.load(input_skl_model)
 
# Convert the sklearn model into ONNX
onnx_model = onnxmltools.convert_sklearn(skl_model, initial_types=input_data_type)
 
# Save as protobuf
onnxmltools.utils.save_model(onnx_model, output_onnx_model)

测试onnx是否正确执行

import onnxruntime as rt
import numpy
from tensorflow.keras.preprocessing import image
import cv2
import numpy as np

# 加载onnx
sess = rt.InferenceSession("model.onnx")
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
probability_name = sess.get_outputs()[1].name

# 输入图片，将其转化成拉普拉斯算子
input_size = (512, 512)
img = image.load_img("./dataset/2.jpg", target_size=input_size)
gray = cv2.cvtColor(np.asarray(img), cv2.COLOR_BGR2GRAY)
laplacian = cv2.Laplacian(gray, cv2.CV_64F)
img_features = [[laplacian.var(), np.amax(laplacian)]]

# 运行模型
pred_onx = sess.run([label_name, probability_name], {input_name: img_features})

msg = '\nModel Prediction: ' + ('Undistorted' if (
    pred_onx[0] == 0) else 'Blurred\n')
print("\033[96m {}\033[00m" .format(msg))

# print info
# print('input_name: ' + input_name)
# print('label_name: ' + label_name)
# print('probability_name: ' + probability_name)
# print(img_features)，features类型[[547.8542016164865, 592.0]]
# print(pred_onx)，pred_onx类型[array([0], dtype=int64), array([[-0.82165027,  0.82165027]], dtype=float32)]，第一个0就是表示非模糊。

参考文档https://blog.csdn.net/qq_35366482/article/details/122683861
官方文档http://onnx.ai/sklearn-onnx/auto_examples/index.html