機(jī)器學(xué)習(xí) | 使用Onnx部署模型
在機(jī)器學(xué)習(xí)中，通常訓(xùn)練完成了模型該如何部署？如果直接將訓(xùn)練環(huán)境部署在線上，會和訓(xùn)練平臺或者框架相關(guān)，所以通常轉(zhuǎn)換為ONNX格式。
1、ONNX簡介
開放神經(jīng)網(wǎng)絡(luò)交換ONNX（Open Neural Network Exchange）是一套表示深度神經(jīng)網(wǎng)絡(luò)模型的開放格式，由微軟和Facebook于2017推出，然后迅速得到了各大廠商和框架的支持。ONNX定義了一組與環(huán)境和平臺無關(guān)的標(biāo)準(zhǔn)格式，為AI模型的互操作性提供了基礎(chǔ)，使AI模型可以在不同框架和環(huán)境下交互使用。硬件和軟件廠商可以基于ONNX標(biāo)準(zhǔn)優(yōu)化模型性能，讓所有兼容ONNX標(biāo)準(zhǔn)的框架受益。目前，ONNX主要關(guān)注在模型預(yù)測方面（inferring），使用不同框架訓(xùn)練的模型，轉(zhuǎn)化為ONNX格式后，可以很容易的部署在兼容ONNX的運(yùn)行環(huán)境中。
2、模型部署工作流程
模型部署工作流程通用如下：
（1）模型訓(xùn)練：使用PyTorch、TensorFlow或其他深度學(xué)習(xí)框架進(jìn)行模型訓(xùn)練；（2）導(dǎo)出為ONNX格式：模型訓(xùn)練完成，將模型導(dǎo)出為ONNX格式；（3）模型轉(zhuǎn)換：使用ONNX轉(zhuǎn)換工具，如ONNX Runtime、TensorRT或OpenVINO等，將ONNX模型轉(zhuǎn)換為目標(biāo)設(shè)備和框架所支持的格式；（4）部署和推理：將轉(zhuǎn)換后的模型部署到目標(biāo)設(shè)備上，并使用相應(yīng)的推理引擎進(jìn)行推理；
部署流程
3、部署模型
為了更直觀的了解ONNX格式內(nèi)容，以下操作一下完整的流程。
3.1 訓(xùn)練模型
這里為了簡單訓(xùn)練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 將訓(xùn)練的模型轉(zhuǎn)換為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 轉(zhuǎn)換是上述代碼，對于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 部署預(yù)測模型
使用 flask 作為serve部署API服務(wù)，加載onnx文件，并處理預(yù)測請求，代碼如下：
# 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)相關(guān)的測試 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]}'
// 返回結(jié)果
{
"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]}'
// 返回結(jié)果
{
"result": [
0
]
}資料
1、https://github.com/aipredict/ai-deployment/blob/master/deploy-ml-dl-using-onnx/README.md
