Torchvision目标检测模型转CoreML：求转换示例及解决方案

Torchvision目标检测模型是可以转换为CoreML的，官方文档未覆盖这类场景是因为检测模型的结构比分类/分割更复杂，需要做额外的适配处理。针对你尝试的ssdlite320_mobilenet_v3_large、fasterrcnn_resnet50_fpn_v2和retinanet_resnet50_fpn_v2，以下是可落地的转换步骤：

核心问题解析 #

你遇到的错误本质是：Torchvision检测模型的前向传播逻辑中会生成ImageList自定义类实例，且包含CoreML转换器无法解析的复杂拆包（prim::TupleUnpack）操作。解决思路是用包装类封装原模型，将输入输出统一为标准张量，避开这些无法识别的结构。

具体转换步骤 #

• 1. 封装模型，统一输入输出为张量
  写一个简单的包装模块，把原模型的输入（图像列表）改成单张量，输出直接提取检测结果的核心字段：

import torch
import torchvision.models.detection as detection

class DetectionWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.model.eval()

    def forward(self, x):
        # x 形状为 [batch, 3, H, W] 的张量
        with torch.no_grad():
            outputs = self.model([x])[0]
        # 提取检测结果的关键输出：边界框、置信度、类别标签
        return outputs["boxes"], outputs["scores"], outputs["labels"]

• 2. 导出为ONNX格式（推荐，比直接转CoreML更稳定）
  以ssdlite320_mobilenet_v3_large为例，其他模型仅需替换模型加载代码：

# 加载预训练检测模型
model = detection.ssdlite320_mobilenet_v3_large(pretrained=True)
wrapper = DetectionWrapper(model)

# 创建示例输入，尺寸需匹配模型要求（SSDLite为320x320，FasterRCNN/RetinaNet支持任意尺寸）
dummy_input = torch.randn(1, 3, 320, 320)

# 导出ONNX模型，配置动态轴支持批量推理
torch.onnx.export(
    wrapper,
    dummy_input,
    "ssdlite.onnx",
    opset_version=13,
    input_names=["image"],
    output_names=["boxes", "scores", "labels"],
    dynamic_axes={
        "image": {0: "batch_size"},
        "boxes": {0: "num_detections"},
        "scores": {0: "num_detections"},
        "labels": {0: "num_detections"}
    }
)

• 3. 转换ONNX到CoreML
  使用coremltools完成最终转换，可将输入设置为图像类型，方便后续直接传入图像：

import coremltools as ct

# 加载ONNX模型规格
onnx_spec = ct.utils.load_spec("ssdlite.onnx")

# 转换为CoreML模型，设置输入图像的缩放系数（匹配PyTorch的图像预处理逻辑）
coreml_model = ct.convert(
    onnx_spec,
    inputs=[ct.ImageType(name="image", shape=dummy_input.shape, scale=1/255.0)]
)

# 保存CoreML模型
coreml_model.save("SSDLiteMobileNet.mlmodel")

适配其他模型 #

对于fasterrcnn_resnet50_fpn_v2和retinanet_resnet50_fpn_v2，仅需修改模型加载代码，包装类和转换流程完全一致。示例输入尺寸可根据需求调整（比如设置为640x640），这两个模型支持动态输入尺寸。

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