机器学习|使用ONNX部署模型

用户1904552

发布于 2025-02-27 10:36:46

32300

代码可运行

文章被收录于专栏：周末程序猿周末程序猿

运行总次数：0

代码可运行

在机器学习中，通常训练完成了模型该如何部署？如果直接将训练环境部署在线上，会和训练平台或者框架相关，所以通常转换为ONNX格式。

1、ONNX简介

开放神经网络交换ONNX（Open Neural Network Exchange）是一套表示深度神经网络模型的开放格式，由微软和Facebook于2017推出，然后迅速得到了各大厂商和框架的支持。 ONNX定义了一组与环境和平台无关的标准格式，为AI模型的互操作性提供了基础，使AI模型可以在不同框架和环境下交互使用。硬件和软件厂商可以基于ONNX标准优化模型性能，让所有兼容ONNX标准的框架受益。目前，ONNX主要关注在模型预测方面（inferring），使用不同框架训练的模型，转化为ONNX格式后，可以很容易的部署在兼容ONNX的运行环境中。

2、模型部署工作流程

模型部署工作流程通用如下：

（1）模型训练：使用PyTorch、TensorFlow或其他深度学习框架进行模型训练；（2）导出为ONNX格式：模型训练完成，将模型导出为ONNX格式；（3）模型转换：使用ONNX转换工具，如ONNX Runtime、TensorRT或OpenVINO等，将ONNX模型转换为目标设备和框架所支持的格式；（4）部署和推理：将转换后的模型部署到目标设备上，并使用相应的推理引擎进行推理；

部署流程

3、部署模型

为了更直观的了解ONNX格式内容，以下操作一下完整的流程。

3.1 训练模型

这里为了简单训练LogisticRegression模型，代码如下：

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression

iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)

clr = LogisticRegression()
clr.fit(X_train, y_train)

3.2 将训练的模型转换为ONNX格式

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

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

这里每个框架不一样，scikit-learn 转换是上述代码，对于PyTorch使用类似如下代码：

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x * x[0].item()
        return x, torch.Tensor([i for i in x])

model = Model()      
dummy_input = torch.rand(10)
torch.onnx.export(model, dummy_input, 'test-pytorch.onnx')

通过上述代码获得 logreg_iris.onnx 文件，加载模型并打印模型信息。

import onnx

model = onnx.load('logreg_iris.onnx')
print(model)

3.3 部署预测模型

使用 flask 作为serve部署API服务，加载onnx文件，并处理预测请求，代码如下：

# python3 -m pip install flask onnxruntime --break-system-packages
from flask import Flask, jsonify, request
import onnxruntime as rt
import numpy as np

app = Flask(__name__)
sess = rt.InferenceSession("logreg_iris.onnx")
# 获取模型输入的名称和形状
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].shape
print(f"input_name: {input_name}, label_name: {label_name}")

@app.route('/api/predict', methods=['POST'])
def predict():
    data = request.json
    if not data:
        return jsonify({'error': 'Invalid input'}), 400
    
    input_data = np.array([data["data"]], dtype=np.float32)
    outputs = sess.run(None, {input_name: input_data})
    print("outputs: ", outputs)
    return jsonify({
        "result": outputs[0].tolist(),
    }), 200

if __name__ == '__main__':
    app.run(debug=True)