利用ONNX Runtime高效运行PyTorch模型
本文将指导您如何使用ONNX Runtime运行经torch.onnx.export导出的PyTorch模型，并重点解决PyTorch张量与ONNX Runtime所需NumPy数组类型不兼容的问题。
首先，我们来看一个PyTorch模型导出示例：
import torch
class SumModule(torch.nn.Module):
def forward(self, x):
return torch.sum(x, dim=1)
torch.onnx.export(
SumModule(),
(torch.ones(2, 2),),
"onnx.pb",
input_names=["x"],
output_names=["sum"]
)
这段代码定义了一个简单的PyTorch模型SumModule，并将其导出为名为onnx.pb的ONNX模型文件。
直接使用PyTorch张量作为ONNX Runtime的输入会导致错误，因为ONNX Runtime期望的是NumPy数组。 错误信息通常提示输入类型错误。
为了解决这个问题，我们需要将PyTorch张量转换为NumPy数组。 正确的代码如下：
import onnxruntime
import numpy as np
import torch
ort_session = onnxruntime.InferenceSession("onnx.pb")
# 关键修改：将torch.Tensor转换为np.ndarray
x = np.ones((2, 2), dtype=np.float32)
inputs = {ort_session.get_inputs()[0].name: x}
print(ort_session.run(None, inputs))
这段代码加载onnx.pb文件，创建一个形状为(2, 2)，数据类型为float32的NumPy数组作为模型输入。 ort_session.get_inputs()[0].name 获取输入张量的名称，确保输入数据与模型定义匹配。 ort_session.run 函数运行模型并打印输出结果。
更简洁的等效代码：
import onnxruntime as ort
import numpy as np
sess = ort.InferenceSession("onnx.pb")
input_data = np.ones((2, 2)).astype(np.float32)
output_data = sess.run(None, {"x": input_data})[0]
print(output_data)
这段代码功能相同，但更简洁易读。 关键在于使用NumPy数组作为输入。
通过以上方法，您可以成功加载并运行使用torch.onnx.export导出的PyTorch模型。 请确保输入数据的类型和形状与模型的预期输入相匹配。
