ONNX Runtime 後端用於 ONNX

ONNX Runtime 擴展了 onnx 後端 API,以使用此執行階段執行預測。讓我們使用 API 來計算簡單邏輯迴歸模型的預測。

import skl2onnx
import onnxruntime
import onnx
import sklearn
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
import numpy
from onnxruntime import get_device
import numpy as np
import onnxruntime.backend as backend

我們先建立 ONNX 圖。

data = load_iris()
X, Y = data.data, data.target
logreg = LogisticRegression(C=1e5).fit(X, Y)
model = skl2onnx.to_onnx(logreg, X.astype(np.float32))
name = "logreg_iris.onnx"
with open(name, "wb") as f:
    f.write(model.SerializeToString())

讓我們使用 ONNX 後端 API 來測試它。

model = onnx.load(name)
rep = backend.prepare(model)
x = np.array(
    [[-1.0, -2.0, 5.0, 6.0], [-1.0, -2.0, -3.0, -4.0], [-1.0, -2.0, 7.0, 8.0]],
    dtype=np.float32,
)
label, proba = rep.run(x)
print("label={}".format(label))
print("probabilities={}".format(proba))
label=[2 0 2]
probabilities=[{0: 0.0, 1: 0.0, 2: 1.0}, {0: 1.0, 1: 1.9515885113950192e-38, 2: 0.0}, {0: 0.0, 1: 0.0, 2: 1.0}]

裝置取決於套件的編譯方式,GPU 或 CPU。

print(get_device())
GPU

後端也可以直接載入模型,而無需使用 onnx

rep = backend.prepare(name)
x = np.array(
    [[-1.0, -2.0, -3.0, -4.0], [-1.0, -2.0, -3.0, -4.0], [-1.0, -2.0, -3.0, -4.0]],
    dtype=np.float32,
)
label, proba = rep.run(x)
print("label={}".format(label))
print("probabilities={}".format(proba))
label=[0 0 0]
probabilities=[{0: 1.0, 1: 1.9515885113950192e-38, 2: 0.0}, {0: 1.0, 1: 1.9515885113950192e-38, 2: 0.0}, {0: 1.0, 1: 1.9515885113950192e-38, 2: 0.0}]

後端 API 由其他框架實作,並且可以更輕鬆地在具有相同 API 的多個執行階段之間切換。

此範例使用的版本

print("numpy:", numpy.__version__)
print("scikit-learn:", sklearn.__version__)
print("onnx: ", onnx.__version__)
print("onnxruntime: ", onnxruntime.__version__)
print("skl2onnx: ", skl2onnx.__version__)
numpy: 1.26.4
scikit-learn: 1.6.dev0
onnx:  1.17.0
onnxruntime:  1.18.0+cu118
skl2onnx:  1.17.0

腳本的總執行時間: (0 分鐘 17.934 秒)

由 Sphinx-Gallery 產生的展示