如何解决使用 ONNX 模型推理时遇到的 "RuntimeError: Input must be a list of dictionaries or a single numpy array" 错误？

珍惜时间，勤奋学习！今天给大家带来《如何解决使用 ONNX 模型推理时遇到的 "RuntimeError: Input must be a list of dictionaries or a single numpy array" 错误？》，正文内容主要涉及到等等，如果你正在学习文章，或者是对文章有疑问，欢迎大家关注我！后面我会持续更新相关内容的，希望都能帮到正在学习的大家！

如何调用 torch.onnx.export 导出的模型？

torch.onnx.export 会创建一个 onnx 模型，该模型可以在其他框架或设备上进行推理。

解决 runtimeerror：输入必须是字典列表或单个 numpy 数组

当你使用导出的 onnx 模型进行推理时，遇到了以下错误：

runtimeerror: input must be a list of dictionaries or a single numpy array for input 'x'.

这是因为 pytorch 模型的输入是张量，而 onnx 模型的输入是数组。要解决此错误，请将输入张量转换为 numpy 数组。

修改后的示例代码：

import onnxruntime
import numpy
import torch

# 创建 ONNX 模型
model = SumModule()
torch.onnx.export(model, (torch.ones(2, 2),), "onnx.pb", input_names=["x"], output_names=["sum"])

# 加载 ONNX 模型
ort_session = onnxruntime.InferenceSession("onnx.pb")

# 为模型创建输入，将其转换为 numpy 数组
input_data = numpy.ones((2, 2), dtype=numpy.float32)

# 运行模型
output_data = ort_session.run(None, {"x": input_data})

# 输出预测结果
print(output_data)

注意：input_data 类型的修改是关键。通过将输入张量转换为 numpy 数组，消除了 pytorch 和 onnx 之间的类型不匹配问题，使推理能够成功进行。

好了，本文到此结束，带大家了解了《如何解决使用 ONNX 模型推理时遇到的 "RuntimeError: Input must be a list of dictionaries or a single numpy array" 错误？》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！