Detectron2模型转ONNX：解决虚拟输入形状错误及最优尺寸获取

错误原因 #

Detectron2的RCNN系列模型（包括Mask R-CNN）的forward方法不直接接收批量张量作为输入，而是要求输入为字典组成的列表。每个字典对应一张图片的输入信息，必须包含"image"键，其值为形状(C, H, W)的单张图片张量。你直接传入了(1,3,224,224)的批量张量，模型尝试将其当作列表遍历并访问x["image"]，自然会触发索引错误。

修正后的导出代码 #

以下是可以正常运行的ONNX导出代码，同时处理了正确的输入格式和输出映射：

import torch
from detectron2.config import get_cfg
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2 import model_zoo

# 配置模型
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 构建并加载模型
model = build_model(cfg)
model.eval()
model.to(cfg.MODEL.DEVICE)
DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)

# 构造正确的虚拟输入：字典组成的列表
# 单张图片输入：形状为(3, H, W)，这里用800x1333（COCO典型尺寸）
dummy_image = torch.randn(3, 800, 1333, device=cfg.MODEL.DEVICE)
dummy_input = [{"image": dummy_image}]

# 定义输入输出名称
input_names = ["input_image"]
# Mask R-CNN的输出包含多个部分，需对应模型的输出结构
output_names = ["pred_boxes", "pred_classes", "pred_scores", "pred_masks"]

# 导出ONNX，开启动态轴以支持可变输入尺寸
torch.onnx.export(
    model,
    dummy_input,
    "mask_rcnn_r50_fpn.onnx",
    verbose=False,
    input_names=input_names,
    output_names=output_names,
    opset_version=11,  # 选择兼容的ONNX版本，建议11+
    dynamic_axes={
        "input_image": {1: "height", 2: "width"},  # 动态高度和宽度
        "pred_boxes": {0: "num_boxes"},
        "pred_classes": {0: "num_boxes"},
        "pred_scores": {0: "num_boxes"},
        "pred_masks": {0: "num_boxes", 2: "mask_height", 3: "mask_width"}
    }
)

print("模型已成功导出为mask_rcnn_r50_fpn.onnx")

最优虚拟输入形状选择 #

1. 参考数据集原生尺寸：COCO数据集的图片尺寸多为800x1333（短边800，长边按比例缩放），使用这个尺寸作为虚拟输入能最大程度匹配模型训练时的输入分布，减少导出后部署的精度偏差。
2. 动态形状优先：如果部署场景需要支持可变尺寸的图片，一定要在导出时设置dynamic_axes，这样ONNX模型可以处理任意(3, H, W)的输入（H、W需满足模型的下采样要求，比如是32的倍数）。
3. 固定尺寸选择：如果必须使用固定尺寸，建议选择常见的检测尺寸，比如640x640、768x1280，确保尺寸是模型下采样步长的整数倍（Mask R-CNN的FPN下采样步长为32，所以H和W最好是32的倍数）。

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