在一個 ONNX 圖表中儲存陣列

一旦模型轉換完成，將陣列作為圖表中的常數儲存並透過輸出擷取會很有用。這讓使用者可以儲存訓練參數或其他資訊，如詞彙表。最後幾節顯示如何移除輸出或將中間結果提升為輸出。

訓練並轉換模型

我們從 :epkg:`ONNX Zoo` 下載一個模型，但該模型可能由其他轉換器函式庫訓練和產生。

import pprint
import numpy
from onnx import load
from onnxruntime import InferenceSession
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from skl2onnx import to_onnx
from skl2onnx.helpers.onnx_helper import (
    add_output_initializer,
    select_model_inputs_outputs,
)


data = load_iris()
X, y = data.data.astype(numpy.float32), data.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
model = LogisticRegression(penalty="elasticnet", C=2.0, solver="saga", l1_ratio=0.5)
model.fit(X_train, y_train)

onx = to_onnx(model, X_train[:1], target_opset=12, options={"zipmap": False})

/home/xadupre/github/scikit-learn/sklearn/linear_model/_sag.py:350: ConvergenceWarning: The max_iter was reached which means the coef_ did not converge
  warnings.warn(

新增訓練參數

new_onx = add_output_initializer(
    onx, ["C", "l1_ratio"], [numpy.array([model.C]), numpy.array([model.l1_ratio])]
)

推論

sess = InferenceSession(new_onx.SerializeToString(), providers=["CPUExecutionProvider"])
print("output names:", [o.name for o in sess.get_outputs()])
res = sess.run(None, {"X": X_test[:2]})
print("outputs")
pprint.pprint(res)

output names: ['label', 'probabilities', 'C', 'l1_ratio']
outputs
[array([1, 0], dtype=int64),
 array([[2.4983161e-03, 8.6063814e-01, 1.3686356e-01],
       [9.7875208e-01, 2.1247936e-02, 2.7418587e-08]], dtype=float32),
 array([2.]),
 array([0.5])]

此解決方案的主要缺點是增加預測時間，因為 onnxruntime 會為每個預測複製常數。可以將這些常數儲存在單獨的 ONNX 圖表中，或將它們移除。

選擇輸出

下一個函式會從模型中移除不需要的輸出，不僅是常數。下一個模型僅保留機率。

simple_onx = select_model_inputs_outputs(new_onx, ["probabilities"])

sess = InferenceSession(
    simple_onx.SerializeToString(), providers=["CPUExecutionProvider"]
)
print("output names:", [o.name for o in sess.get_outputs()])
res = sess.run(None, {"X": X_test[:2]})
print("outputs")
pprint.pprint(res)

# Function *select_model_inputs_outputs* add also promote an intermediate
# result to an output.
#

output names: ['probabilities']
outputs
[array([[2.4983161e-03, 8.6063814e-01, 1.3686356e-01],
       [9.7875208e-01, 2.1247936e-02, 2.7418587e-08]], dtype=float32)]

此範例僅使用記憶體中的 ONNX 圖表，從不儲存或載入模型。這可以使用以下程式碼片段完成。

儲存模型

with open("simplified_model.onnx", "wb") as f:
    f.write(simple_onx.SerializeToString())

載入模型

model = load("simplified_model.onnx")

sess = InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
print("output names:", [o.name for o in sess.get_outputs()])
res = sess.run(None, {"X": X_test[:2]})
print("outputs")
pprint.pprint(res)

output names: ['probabilities']
outputs
[array([[2.4983161e-03, 8.6063814e-01, 1.3686356e-01],
       [9.7875208e-01, 2.1247936e-02, 2.7418587e-08]], dtype=float32)]