为ONNX格式Complex-YOLO模型添加NMS层报错及张量裁剪求助

问题背景 #

导出的Complex-YOLO ONNX模型输出张量shape为[num_batches, spatial_dimension, 10]，添加NMS层时因NMS要求boxes输入的最后一维必须是4，运行时报错：

non_max_suppression.cc:92 PrepareCompute The most inner dimension in boxes must have 4 data.

解决思路 #

通过ONNX的Slice节点从模型输出张量中提取最后一维的前4个元素，得到符合NMS要求的boxes张量（shape为[num_batches, spatial_dimension,4]）。

修改后的完整代码 #

在原代码中新增Slice节点提取boxes，并更新NMS的输入：

import onnx
from onnx import TensorProto

onnx_model = onnx.load("Complex_yolo_epoch_300.onnx")
onnx.checker.check_model(onnx_model)
onnx_model_with_nms_path = "Complex_yolo_epoch_300_with_nms.onnx"

graph = onnx_model.graph

# -------------------------- 新增：添加Slice节点提取前4维度作为boxes --------------------------
# 定义Slice的起始、结束、步长张量
start = onnx.helper.make_tensor("start", TensorProto.INT64, [3], [0, 0, 0])
end = onnx.helper.make_tensor("end", TensorProto.INT64, [3], [0, 0, 4])
step = onnx.helper.make_tensor("step", TensorProto.INT64, [3], [1, 1, 1])

# 创建Slice节点，从模型输出中提取前4个通道作为boxes
slice_boxes_node = onnx.helper.make_node(
    'Slice',
    inputs=['output', 'start', 'end', 'step'],
    outputs=['boxes'],
)
graph.node.append(slice_boxes_node)

# 将Slice的参数张量添加到模型初始化器
graph.initializer.append(start)
graph.initializer.append(end)
graph.initializer.append(step)
# -------------------------------------------------------------------------------------------

# 转置分数张量（原代码保留）
transpose_scores_node = onnx.helper.make_node(
    'Transpose', 
    inputs=['output'], 
    outputs=['scores_transposed'],
    perm=(0, 2, 1))
graph.node.append(transpose_scores_node)

# 创建NMS所需的常量张量（原代码保留）
score_threshold = onnx.helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.25])
iou_threshold = onnx.helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.45])
max_output_boxes_per_class = onnx.helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [200])

# -------------------------- 修改：NMS输入使用提取后的boxes --------------------------
inputs=["boxes", 'scores_transposed', 'max_output_boxes_per_class', 'iou_threshold', 'score_threshold']
# -----------------------------------------------------------------------------------

# 创建NMS节点（原代码保留）
nms_node = onnx.helper.make_node(
    'NonMaxSuppression',
    inputs,
    outputs=["selected_indices"],
    center_point_box=1, 
)
graph.node.append(nms_node)

# 添加输出节点信息（原代码保留）
output_value_info = onnx.helper.make_tensor_value_info("selected_indices", TensorProto.INT64, shape=["num_results",3])
graph.output.append(output_value_info)

# 添加常量初始化器（原代码保留）
graph.initializer.append(score_threshold)
graph.initializer.append(iou_threshold)
graph.initializer.append(max_output_boxes_per_class)

# 检查并保存模型（原代码保留）
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, onnx_model_with_nms_path)

代码说明 #

• Slice节点的end参数前两个维度设为0，用于保留输入张量的动态维度（batch_size和spatial_dimension），仅限制最后一维取前4个元素。
• 修改NMS节点的输入为提取后的boxes张量，确保符合NMS对输入shape的要求。
• 所有自定义常量张量（start、end、step等）必须添加到graph.initializer，否则ONNX Runtime会报错找不到这些张量。

内容的提问来源于stack exchange，提问作者user1080814