Detectron2模型转ONNX时触发IndexError错误求助

错误原因 #

报错IndexError: too many indices for tensor of dimension 2的核心原因：

• Detectron2模型要求输入是字典组成的列表（即batched_inputs），每个字典包含"image"键，对应单张图片的3维张量（C,H,W）；你直接传入单个3维张量，导致模型内部遍历输入时解析出错。
• 未将模型切换到eval模式，导出ONNX前必须关闭BatchNorm、Dropout等训练相关层的影响。
• 输入缺少batch维度，ONNX通常期望带batch维度的张量（1,C,H,W）。

修正后的完整代码 #

import torch
import onnx
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.modeling import build_model
from detectron2.data.transforms import ResizeShortestEdge as T
from detectron2.export import export_onnx_model

# 1. 配置并加载预训练模型
cfg = get_cfg()
cfg.MODEL.DEVICE = 'cpu'
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
# 加载官方预训练权重，必须添加此步骤，否则导出的是随机初始化模型
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
model = build_model(cfg)
# 切换模型到推理模式
model.eval()

# 2. 预处理输入图像（假设im是你的numpy格式原图，shape为(H,W,C)）
aug = T([cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST)
image = aug.get_transform(im).apply_image(im)
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
# 增加batch维度，变为(1,3,H,W)
image = image.unsqueeze(0)
# 包装成模型要求的输入格式：字典组成的列表
batched_inputs = [{"image": image}]

# 3. 导出并验证ONNX模型
output_path = "/home/ecoation/Documents/model/deploy.onnx"
export_onnx_model(
    model,
    batched_inputs,
    output_path,
    opset_version=11,  # 选择兼容目标ONNX Runtime的opset版本，11为通用兼容版本
    do_constant_folding=True
)

# 验证导出的模型有效性
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型导出成功且验证通过")

关键修正点说明 #

• 加载预训练权重：原代码未指定预训练权重，导出的是随机初始化模型，必须通过cfg.MODEL.WEIGHTS加载官方预训练权重。
• 模型切换到eval模式：model.eval()是导出ONNX的必要步骤，确保模型处于推理状态。
• 修正输入格式：将图像张量包装成[{"image": 带batch维度的张量}]的格式，完全匹配Detectron2模型的输入要求。
• 添加batch维度：通过unsqueeze(0)给图像张量增加batch维度，满足ONNX的输入规范。

内容的提问来源于stack exchange，提问作者Frank van Paassen