如何在Python上运行ONNX模型文件

Question

我有一个ONNX模型文件，其中包含文本分类器。我也有数据，我的目标是在新的数据上测试模型。但是我无法运行这些代码。首先我尝试下面的代码，什么都没发生

model_name = 'text_model.onnx'
onnx_model = onnx.load(model_name)
onnx.checker.check_model(onnx_model)

然后我试试这个

ort.InferenceSession(onnx_model, None)

但是我得到一个错误：

TypeError: Unable to load from type '<class 'onnx.onnx_ml_pb2.ModelProto'>'

我该怎么办？

Answer 1

对我有效的是以下代码

from onnxruntime import InferenceSession
filename = "./svm_iris.onnx"
def execute_onnx_model_from_file(filename: str) -> None:
sess = InferenceSession(filename)
x_test, y_test = json_to_ndarray()
sess.run(None, {"X": x_test.astype(np.float32)})[0]

在文档中，指定您可以将存储的ONNX模型的文件名交给推理会话。另一个更接近代码的解决方案是序列化onnx模型：

from onnxruntime import InferenceSession
filename = "./svm_iris.onnx"
def execute_onnx_model_from_file(filename: str) -> None:
onnx_model = onnx.load(filename)
sess = InferenceSession(onnx_model.SerializeToString())
x_test, y_test = json_to_ndarray()
sess.run(None, {"X": x_test.astype(np.float32)})[0]