以下是一个完整的示例,展示如何加载 ONNX 模型、获取模型信息并运行推理:
import onnxruntime
import numpy as np
# 模型路径
bev_head_onnx_path = "path/to/your/bev_head.onnx"
# 加载模型
session = onnxruntime.InferenceSession(bev_head_onnx_path)
# 获取模型元信息
model_meta = session.get_modelmeta()
print("Model Name:", model_meta.name)
print("Model Version:", model_meta.version)
print("Inputs:", [input.name for input in model_meta.graph_inputs])
print("Outputs:", [output.name for output in model_meta.graph_outputs])
# 准备输入数据
# 假设模型的输入名称为 'volumes',形状为 (1, 3, 224, 224)
volumes = np.random.rand(1, 3, 224, 224).astype(np.float32)
# 运行推理
onnx_results = session.run([], {'volumes': volumes})
# 打印输出结果
print("Output shape:", [result.shape for result in onnx_results])
输出示例
假设模型的输入名称为 'volumes'
,输出名称为 'output'
,运行上述代码可能会输出类似以下内容:
Model Name: bev_head_model
Model Version: 1.0
Inputs: ['volumes']
Outputs: ['output']
Output shape: [(1, 1000)]
注意事项
-
输入名称和形状:
-
输入数据的名称(如
'volumes'
)和形状必须与 ONNX 模型中定义的输入一致。可以通过get_modelmeta()
检查输入名称和形状。 -
如果输入数据的形状不正确,可能会导致运行时错误。
-
-
数据类型:
- 输入数据的类型(如
np.float32
)也必须与模型的要求一致。可以通过model_meta.graph_inputs
检查输入数据类型。
- 输入数据的类型(如
-
输出结果:
session.run()
返回的结果是一个列表,其中每个元素对应一个输出张量。如果模型有多个输出,可以通过指定输出名称列表来选择需要的输出。
-
性能优化:
- 如果需要更高的推理性能,可以使用
onnxruntime
的优化选项,例如设置执行提供者(如CUDAExecutionProvider
)或启用图优化。
- 如果需要更高的推理性能,可以使用
通过这种方式,你可以方便地加载和运行 ONNX 模型,并获取模型的推理结果。