教程#

ONNX Runtime 提供了一種簡單的方法，可以在 CPU 或 GPU 上高效能執行機器學習模型，且無需依賴於訓練框架。機器學習框架通常針對批次訓練而非預測進行最佳化，而預測是應用程式、網站和服務中更常見的場景。在宏觀層面，您可以：

  1. 使用您喜歡的框架訓練模型。

  2. 將模型轉換或匯出為 ONNX 格式。詳情請參閱 ONNX 教程。

  3. 使用 ONNX Runtime 載入並執行模型。

在本教程中，我們將簡要地使用 scikit-learn 建立一個管道，將其轉換為 ONNX 格式並執行首次預測。

步驟 1：使用您喜歡的框架訓練模型#

我們將使用著名的鳶尾花資料集。

<<<

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)

from sklearn.linear_model import LogisticRegression
clr = LogisticRegression()
clr.fit(X_train, y_train)
print(clr)

>>>

    LogisticRegression()

步驟 2：將模型轉換或匯出為 ONNX 格式#

ONNX 是一種描述機器學習模型的格式。它定義了一組常用的運算子來組合模型。有工具可以將其他模型格式轉換為 ONNX。這裡我們將使用 ONNXMLTools。

<<<

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

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())

>>>

步驟 3：使用 ONNX Runtime 載入並執行模型#

我們將使用 ONNX Runtime 計算此機器學習模型的預測。

<<<

import numpy
import onnxruntime as rt

sess = rt.InferenceSession(
    "logreg_iris.onnx", providers=rt.get_available_providers())
input_name = sess.get_inputs()[0].name
pred_onx = sess.run(None, {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)

>>>

    [1 2 1 2 2 2 0 1 0 2 2 2 2 1 1 1 1 2 0 1 0 2 0 2 0 0 0 0 2 2 1 2 1 1 1 0 2
     1]

可以透過將特定輸出的名稱指定到列表中來更改程式碼以獲取該特定輸出。

<<<

import numpy
import onnxruntime as rt

sess = rt.InferenceSession(
    "logreg_iris.onnx", providers=rt.get_available_providers())
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)

>>>

    [1 2 1 2 2 2 0 1 0 2 2 2 2 1 1 1 1 2 0 1 0 2 0 2 0 0 0 0 2 2 1 2 1 1 1 0 2
     1]