簡介

快速入門

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

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

  2. 將模型轉換或匯出為 ONNX 格式。有關更多詳細資訊，請參閱ONNX 教學。

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

在本教學中，我們將簡要建立一個具有 scikit-learn 的管線，將其轉換為 ONNX 格式並執行第一個預測。

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

我們將使用著名的 Iris 資料集。

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)

步驟 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 onnxruntime as rt
sess = rt.InferenceSession("logreg_iris.onnx", providers=["CPUExecutionProvider"])
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]

convert_sklearn, to_onnx, initial_types

該模組實作了兩個函式：convert_sklearn 和 to_onnx。第一個函式已在先前的範例中使用，它需要兩個強制參數

  • 一個 scikit-learn 模型或管線

  • 初始類型

scikit-learn 不會儲存有關訓練資料集的資訊。並非總是能夠檢索特徵的數量或其類型。這就是為什麼函式需要另一個名為 initial_types 的參數的原因。在許多情況下，訓練資料集是一個數值矩陣 X_train。然後它變成 initial_type=[('X', FloatTensorType([None, X_train.shape[1]]))]。X 是這個唯一輸入的名稱，第二個詞表示類型和形狀。形狀是 [None, X_train.shape[1]]，第一個維度是行數，後跟特徵數。行數未定義，因為在轉換模型時，請求的預測數量是未知的。特徵數量通常是已知的。現在假設輸入是一個字串欄，後跟一個矩陣，那麼初始類型將是

initial_type=[
    ('S', StringTensorType([None, 1])),
    ('X', FloatTensorType([None, X_train.shape[1]])),
]

函式 to_onnx 是在與 scikit-learn 的核心開發人員討論後實作的。它還包含一種機制，可以根據訓練資料集的一行推斷出適當的類型。然後，以下程式碼 convert_sklearn(clr, initial_types=[('X', FloatTensorType([None, 4]))]) 通常會重寫為 to_onnx(clr, X_train[:1])，其中 X_train 是訓練資料集，它可以是矩陣或資料框架。預設情況下，輸入名稱為 'X'，除非 X_train 是資料框架。在這種情況下，欄名稱會用作輸入名稱。