機器學習 | 使用Onnx部署模型
在機器學習中，通常訓練完成了模型該如何部署？如果直接將訓練環境部署在線上，會和訓練平臺或者框架相關，所以通常轉換為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)相關的測試 curl 請求如下：
// 請求
curl -i 'http://127.0.0.1:5000/api/predict' -H 'Content-Type: application/json' -d '{"data":[1,2.6,4.4,1.2]}'
// 返回結果
{
"result": [
1
]
}
// 請求
curl -i 'http://127.0.0.1:5000/api/predict' -H 'Content-Type: application/json' -d '{"data":[1,20000000000,4.4,1.2]}'
// 返回結果
{
"result": [
0
]
}資料
1、https://github.com/aipredict/ai-deployment/blob/master/deploy-ml-dl-using-onnx/README.md
