使用coremltools将Detectron2 TorchScript转CoreML时遇KeyError

这个错误源于coremltools对Detectron2导出的TorchScript模型中ImageList自定义结构的支持不完善，尤其是脚本化模型（scripted model）的转换属于实验性功能，更容易触发这类问题。以下是可行的解决方案：

1. 优先使用Torch追踪模型（Traced Model） #

coremltools明确提示脚本模型支持是实验性的，改用追踪模型能避开大部分自定义结构的兼容性问题。修改导出逻辑如下：

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

def export_traced_model():
    # 加载你的模型配置和权重
    cfg = get_cfg()
    cfg.merge_from_file("your_model_config.yaml")
    cfg.MODEL.WEIGHTS = "your_trained_weights.pth"
    cfg.MODEL.DEVICE = "cpu"  # 必须用CPU环境追踪模型
    model = build_model(cfg)
    model.eval()
    DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)

    # 创建匹配模型输入尺寸的张量作为示例输入
    dummy_input = torch.randn(1, 3, 224, 224)
    # 执行模型追踪
    traced_model = torch.jit.trace(model, (dummy_input,))
    traced_model.save("traced_detectron_model.pt")
    return traced_model

将生成的traced_model传入你的CoreML转换函数即可，原转换代码无需大幅修改。

2. 用包装类屏蔽ImageList依赖（适用于必须用脚本模型的场景） #

如果必须使用脚本化模型，可以给原模型套一层包装类，手动处理ImageList的创建逻辑，让模型直接以张量作为输入输出：

class DetectronWrapper(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.model = original_model

    def forward(self, image_tensor):
        from detectron2.structures import ImageList
        # 手动将输入张量转换为模型需要的ImageList格式
        image_list = ImageList.from_tensors([image_tensor], self.model.backbone.size_divisibility)
        outputs = self.model(image_list)
        # 将输出转换为CoreML可识别的纯张量格式（根据你的任务调整返回内容）
        return outputs["instances"].scores, outputs["instances"].pred_boxes.tensor

对这个包装类执行脚本化或追踪后，再传入coremltools进行转换。

3. 升级依赖版本提升兼容性 #

确保Detectron2使用最新稳定版，同时将coremltools升级到v7及以上版本，新版本对Detectron2这类复杂模型的转换支持更完善。

转换完成后，可通过Xcode直接打开生成的.mlmodel文件检查输入输出结构，或用coremltools工具验证推理结果：

import coremltools as ct
model = ct.models.MLModel("newmodel.mlmodel")
# 用示例图片测试
import PIL.Image
img = PIL.Image.open("test_image.jpg").resize((224,224))
pred = model.predict({"image": img})
print(pred)

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