如何使用 torch.onnx.export 导出的 ONNX 模型进行预测？

调用 torch.onnx.export 导出的 onnx 模型

本文旨在解答如何使用 torch.onnx.export 导出的 onnx 模型。

问题：

如何使用 torch.onnx.export 导出的 onnx.pb 模型文件？

解答：

pytorch 模型的输入为 tensor，而 onnx 的输入为 numpy 数组。因此，将输入数据从 tensor 转换为 numpy 数组即可解决问题。

示例代码：

import onnxruntime
import numpy as np

resnet_onnx = onnxruntime.InferenceSession('onnx.pb')
x = np.ones((2, 2), dtype=np.float32)
inputs = {resnet_onnx.get_inputs()[0].name: x}
print(resnet_onnx.run(None, inputs))

注意：

• 将 torch.tensor 转换为 numpy 数组。
• 使用 numpy 数组作为 onnx 模型的输入。
• onnx 模型的输出是一个列表，包含 numpy 数组。

以上就是如何使用 torch.onnx.export 导出的 ONNX 模型进行预测？的详细内容，更多请关注硕下网其它相关文章！