pytorch模型导出为onnx并进行推理
本文将详细介绍如何使用torch.onnx.export导出pytorch模型到onnx格式，以及如何使用onnx runtime进行推理。
许多开发者在使用pytorch训练模型后，希望能够将其部署到其他环境中，例如移动端或服务器端。这时，onnx (open neural network exchange) 格式就发挥了重要作用。onnx 是一种开放的标准，允许不同框架之间互操作。torch.onnx.export 函数可以将pytorch模型导出为onnx格式。然而，导出后的onnx模型如何使用呢？
以下代码展示了使用torch.onnx.export导出一个简单的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"]
)这段代码导出一个名为onnx.pb的onnx模型文件。 该模型接收一个形状为(2, 2)的张量作为输入，并计算其每一行的和作为输出。
接下来，我们需要使用onnx runtime加载并运行这个模型。 直接使用pytorch的tensor作为onnx runtime的输入是错误的。onnx runtime 期望的是numpy数组。
错误的代码示例：
import onnxruntime
import torch
resnet_onnx = onnxruntime.inferencesession("onnx.pb")
x = torch.ones(2, 2)
inputs = {resnet_onnx.get_inputs()[0].name: x}
resnet_onnx.run(none, inputs)这段代码会抛出runtimeerror: input must be a list of dictionaries or a single numpy array for input 'x'的错误。
正确的代码示例：
import onnxruntime
import numpy
import torch
resnet_onnx = onnxruntime.inferencesession("onnx.pb")
x = numpy.ones((2, 2), dtype=numpy.float32)
inputs = {resnet_onnx.get_inputs()[0].name: x}
print(resnet_onnx.run(none, inputs))通过将输入数据类型从torch.tensor更改为numpy.ndarray，并指定数据类型为numpy.float32，我们解决了这个问题。 numpy.ones((2, 2), dtype=numpy.float32) 创建了一个形状为(2, 2)的numpy数组，其元素都为1.0，且数据类型为32位浮点数。 这与onnx模型的预期输入相匹配。 这使得我们可以成功运行onnx模型并得到结果。
另一种使用onnx runtime加载并运行模型的方法：
import onnxruntime as ort
import numpy as np
sess = ort.InferenceSession("onnx.pb")
input_data = np.ones((2, 2)).astype("float32")
output_data = sess.run(None, {"x": input_data})[0]
print(output_data)此方法同样使用了numpy数组作为输入，并成功运行了模型。 需要注意的是，输入数据的形状和数据类型必须与导出onnx模型时指定的输入相匹配。
