torch.onnx.export()报IndexError:二维张量索引过多问题排查

问题复现 #

尝试将Detectron2模型导出为ONNX格式时触发运行报错，使用的导出代码如下：

dummy_input = torch.randn(800, 1067, 3)
torch.onnx.export(model, dummy_input, "/content/model/detectron2.onnx", opset_version=11)

运行后抛出错误栈核心信息如下：

22 dummy_input = torch.randn(800, 1067, 3)
---> 23 torch.onnx.export(model, dummy_input, "/content/model/detectron2.onnx", opset_version=11)
     24 
     25 # Export to Caffe2

15 frames
/usr/local/lib/python3.7/dist-packages/detectron2/modeling/meta_arch/rcnn.py in <listcomp>(.0)
    226         Normalize, pad and batch the input images.
    227         """
---> 228         images = [self._move_to_current_device(x["image"]) for x in batched_inputs]
    229         images = [(x - self.pixel_mean) / self.pixel_std for x in images]
    230         images = ImageList.from_tensors(

IndexError: too many indices for tensor of dimension 2

错误最终定位到Detectron2 RCNN元架构的输入预处理逻辑，提示维度为2的张量存在过多索引。

报错根因 #

错误本质是输入格式、导出方式和模型预期不匹配导致，核心问题有三个：

• Detectron2的RCNN类模型原生forward接口的入参规范是字典组成的列表，每个字典对应一张输入图片，必须包含image键存储对应图片张量。直接传入shape为(800, 1067, 3)的裸张量时，模型内部遍历入参会把第一维长度为800的每个元素当成单个样本字典，尝试取x["image"]，但此时单个元素是shape为(1067, 3)的二维张量，按字典方式索引自然触发维度错误。
• 构造的dummy输入维度顺序不符合PyTorch视觉模型规范：PyTorch CV模型默认输入是(batch, 通道, 高, 宽)的CHW格式，上述代码构造的是HWC格式，即使绕过输入校验也会出现计算逻辑错误。
• Detectron2原生forward中包含大量列表、字典、动态padding的Python逻辑，不属于静态张量计算图，直接调用torch.onnx.export就算修好输入维度也大概率导出失败。

修复方案 #

按优先级选择以下方案即可：

• 优先使用Detectron2官方提供的ONNX导出适配工具，不要直接裸调torch.onnx.export。官方工具已经提前替换了模型中不兼容ONNX的预处理、后处理逻辑，自动处理输入格式、算子兼容问题，是导出成功率最高的方式。
• 如果需要手动导出，先为模型编写一层导出专用的包装类，把原生的字典输入逻辑封装到包装层内部：
  • 包装层的forward直接接收标准(B, C, H, W)格式的四维张量
  • 在包装层内部完成像素均值方差归一化、维度转换逻辑，绕过原生RCNN的输入预处理分支
  • 构造匹配包装层入参格式的dummy输入，正确写法如下：

    # 单batch、3通道、高800、宽1067，同步迁移到模型所在设备
    dummy_input = torch.randn(1, 3, 800, 1067).to(next(model.parameters()).device)
    

• 导出时将opset版本调高到16及以上，opset11版本过低，不支持Detectron2用到的大量新算子，会触发算子不兼容报错。

提示：如果要导出包含检测后处理（NMS、框坐标还原）的全模型，不建议直接追踪原生模型——这部分逻辑是动态Python控制流，无法被ONNX正确序列化，建议只导出核心的特征提取+检测头计算部分，后处理用ONNX算子或者外部逻辑单独实现。

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