PyTorch EfficientNet模型转CoreML时TorchScript报错如何解决

是否所有PyTorch模型都能转换为CoreML模型？ #

不是。CoreML对PyTorch模型的转换依赖两层支持：一是PyTorch自身的TorchScript需要支持模型中用到的所有操作，二是coremltools需要实现对应PyTorch算子的转换逻辑。只要模型中有自定义Python操作、动态逻辑、TorchScript未覆盖的内置方法，就会出现转换失败的情况，你遇到的pythonop不支持、apply方法不兼容都属于这类问题。

无需深入研究算子的解决方案 #

• 方案1：替换第三方EfficientNet实现为TorchVision官方版本
  你当前使用的efficientnet_pytorch第三方库的实现包含TorchScript不兼容的逻辑，直接替换为PyTorch官方TorchVision内置的EfficientNet实现即可，官方实现全量适配TorchScript，仅需调整权重的key名称适配即可，修改后的加载代码如下：

import torch
import torch.nn as nn
import coremltools as ct
from torchvision.models import efficientnet_b2, EfficientNet_B2_Weights

# 加载官方实现的EfficientNetB2
model = efficientnet_b2(weights=EfficientNet_B2_Weights.DEFAULT)
# 替换分类头匹配你的308分类
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, 308)
# 加载预训练权重，替换key名称适配官方实现的层名
prev_state = torch.load('food308_efnetb2_91.31.pth', map_location=torch.device('cpu'))
adjusted_state = {}
for k, v in prev_state.items():
    # 将第三方实现的_fc层名替换为官方实现的classifier.1
    if k.startswith('_fc'):
        k = k.replace('_fc', 'classifier.1')
    adjusted_state[k] = v
model.load_state_dict(adjusted_state)
model.eval()

# 统一输入尺寸进行转换
example_input = torch.rand(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input", shape=example_input.shape)],
)

• 方案2：通过ONNX作为中间格式绕开TorchScript限制
  如果不想修改模型实现，可以先将模型导出为ONNX格式，再从ONNX转换为CoreML，该路径不需要依赖TorchScript的算子支持，适配性更强，转换代码示例如下：

# 接你原有的模型加载代码，导出ONNX
example_input = torch.rand(1, 3, 224, 224)
torch.onnx.export(
    model,
    example_input,
    "efficientnet.onnx",
    export_params=True,
    opset_version=15,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=None
)
# 从ONNX转换为CoreML
mlmodel = ct.convert("efficientnet.onnx", inputs=[ct.TensorType(name="input", shape=example_input.shape)])

• 额外注意：你原有代码中存在输入尺寸不匹配的问题，tracing时用的是(1,3,224,224)，转CoreML时写的是(1,3,64,64)，尺寸不一致也会导致转换或推理出错，转换时请保持输入尺寸统一。

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