PyTorch模型转换为Core ML时torch.jit.trace报错如何解决？

• 最常见的原因是你直接用torch.load()加载的model_final.pth仅存储了模型权重字典（state_dict），并非完整的可执行模型实例，torch.jit.trace无法对字典类型执行追踪。
• 若你存储的是完整模型，可能是模型内部存在动态控制流（如条件判断、循环、动态尺寸计算等），torch.jit.trace仅能记录固定输入对应的执行路径，遇到动态逻辑就会触发报错。
• 输入预处理不匹配、模型与输入张量的设备不统一（比如模型在GPU、输入在CPU）、输入尺寸/通道数不符合模型要求也会导致该报错。

1. 正确加载模型 #

大部分训练场景下导出的.pth文件仅存权重，你需要先实例化训练时定义的模型类，再加载权重参数，同时切换到评估模式：

# 导入你训练时定义的模型类，替换为你实际的模型类路径
from your_model_file import YourModelClass

# 实例化模型
model = YourModelClass()
# 加载权重，指定map_location='cpu'避免设备不匹配问题
model.load_state_dict(torch.load('/content/drive/MyDrive/model/model_final.pth', map_location='cpu'))
# 切换评估模式，关闭训练专属的dropout、batchnorm更新逻辑
model.eval()

2. 对齐输入预处理 #

输入预处理必须和训练阶段完全一致，补充训练时用到的Resize、Normalize等操作，同时确保输入通道、设备和模型匹配：

from torchvision import transforms
from PIL import Image

# 预处理配置和训练时完全对齐，这里以常用的ImageNet预处理为例，替换为你自己的配置
preprocess = transforms.Compose([
    transforms.Resize((224, 224)), # 训练时的输入尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

input_image = Image.open("/content/drive/MyDrive/model/g0079.jpg").convert('RGB') # 强制转3通道RGB，避免灰度图、带alpha通道的图报错
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # 增加batch维度

3. 完成TorchScript转换 #

如果模型存在动态逻辑，优先使用torch.jit.script替代trace，转换完成后验证输出一致性：

# 优先尝试trace，存在动态逻辑时切换为script
try:
    jit_model = torch.jit.trace(model, input_batch)
except:
    jit_model = torch.jit.script(model)

# 验证转换前后输出一致，误差在允许范围内则转换成功
with torch.no_grad():
    origin_out = model(input_batch)
    jit_out = jit_model(input_batch)
print(torch.allclose(origin_out, jit_out, atol=1e-5))

4. 转换为Core ML模型 #

使用coremltools完成最终转换，可直接配置图像输入适配Core ML的图像读取逻辑：

import coremltools as ct

# 转换模型，ImageType配置会自动对齐Core ML的图像输入格式，无需额外处理归一化
coreml_model = ct.convert(
    jit_model,
    inputs=[ct.ImageType(
        name="image_input",
        shape=input_batch.shape,
        scale=1/(0.229*255),
        bias=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
        color_layout=ct.colorlayout.RGB
    )],
    convert_to="neuralnetwork" # 如需兼容新Core ML特性可改为"mlprogram"
)

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

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