如何将PyTorch模型导出为ONNX并使用ONNX Runtime进行推理？

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数组。

错误的代码示例：

盘古大模型

华为云推出的一系列高性能人工智能大模型

35

查看详情

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模型时指定的输入相匹配。

以上就是如何将PyTorch模型导出为ONNX并使用ONNX Runtime进行推理？的详细内容，更多请关注php中文网其它相关文章！