3、模型转换基础:PyTorch模型转ONNX、ONNX转Core ML (.mlmodel)

模型转换这件事,说白了就是让模型换个“语言”说话。PyTorch训练出来的模型,iOS不认识。你得先把它翻译成ONNX这个中间语言,再转成Core ML的.mlmodel格式。嗯,这一步卡住了不少人。

3.1 为什么非要走ONNX这一步?

你可能会问:为什么不直接从PyTorch转Core ML?

我个人习惯走ONNX这条路,原因很简单——ONNX是工业界的“通用语”。PyTorch直接转Core ML的工具链,说实话,坑比较多。我遇到过几次转换失败的情况,最后都是通过ONNX绕过去的。

ONNX(Open Neural Network Exchange)相当于一个中间桥梁。它把PyTorch的模型图、权重、算子全部标准化。Core ML的转换器对ONNX的支持非常成熟,踩坑概率低很多。

核心流程:

PyTorch模型 → ONNX模型 → Core ML模型 (.mlmodel)

3.2 第一步:PyTorch转ONNX

这一步其实不复杂。PyTorch官方提供了torch.onnx.export接口。我直接上代码,你跟着走一遍就明白了。

import torch
import torchvision.models as models

# 加载预训练模型
model = models.mobilenet_v2(pretrained=True)
model.eval()

# 创建一个示例输入(batch_size=1, 3通道, 224x224)
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 示例输入
    "mobilenet_v2.onnx", # 输出文件名
    export_params=True,  # 导出训练好的参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

print("ONNX导出成功!")

这里有几个关键点,我特别提醒一下:

  • opset_version:我建议用11或12。太低的版本不支持某些算子,太高的版本Core ML可能不兼容。
  • dynamic_axes:这个参数很实用。它允许你导出动态batch的模型。我在项目中经常需要处理不同大小的输入,这个参数帮我省了不少事。
  • do_constant_folding:建议开启。它能提前计算一些常量表达式,减小模型体积。

注意:导出前一定要调用model.eval()。我见过有人忘记这步,结果导出的ONNX里包含了Dropout和BatchNorm的训练状态,推理结果完全不对。

3.3 验证ONNX模型

导出之后,我习惯用ONNX Runtime验证一下,确保模型没坏。

import onnx
import onnxruntime as ort
import numpy as np

# 检查模型结构
onnx_model = onnx.load("mobilenet_v2.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过!")

# 用ONNX Runtime跑一次推理
ort_session = ort.InferenceSession("mobilenet_v2.onnx")
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = ort_session.run([output_name], {input_name: test_input})
print(f"推理输出shape: {outputs[0].shape}")

这一步其实很关键。我曾经有一次导出ONNX时忘了设置opset_version,结果模型结构检查通过了,但推理结果全是NaN。后来用ONNX Runtime一跑才发现问题。

3.4 第二步:ONNX转Core ML

这一步我用Apple官方的coremltools。它支持直接从ONNX转换,非常方便。

import coremltools as ct

# 加载ONNX模型
onnx_model = ct.converters.onnx.convert(
    model="mobilenet_v2.onnx",
    minimum_ios_deployment_target='13'
)

# 设置模型元数据
onnx_model.author = 'Your Name'
onnx_model.short_description = 'MobileNetV2 for iOS'
onnx_model.version = '1.0'

# 保存为.mlmodel
onnx_model.save("MobileNetV2.mlmodel")
print("Core ML模型导出成功!")

这里有个细节:minimum_ios_deployment_target参数。我建议至少设到iOS 13。为什么?因为iOS 13开始,Core ML对ONNX的算子支持更完整。如果你设得太低,转换时可能会报错。

小技巧:转换完成后,可以用Xcode打开.mlmodel文件。Xcode会自动生成模型预览,包括输入输出格式、模型大小、甚至预估的推理性能。我每次都会看一眼,确认没问题再集成到项目里。

3.5 常见问题与避坑指南

这条路我走过很多遍,踩过的坑也不少。下面这几个问题,你大概率会遇到:

问题 原因 解决方案
ONNX导出报错:Unsupported operator PyTorch用了ONNX不支持的算子 升级opset_version,或替换为支持的算子
Core ML转换失败:Shape mismatch ONNX模型中有动态shape 在导出ONNX时固定输入shape,或使用dynamic_axes
推理结果与PyTorch不一致 模型没有切换到eval模式 导出前调用model.eval()
模型体积过大 权重没有量化 转换Core ML时使用fp16量化

我曾经遇到过一个特别头疼的问题:MobileViT模型转ONNX时,有个叫torch.nn.functional.interpolate的算子报错。后来发现是opset_version设得太低,升级到12就解决了。嗯,这种问题其实很常见,别慌。

3.6 量化:让模型更轻量

如果你觉得模型太大,可以考虑量化。Core ML支持fp16和int8量化。我一般先用fp16,效果不错,模型体积直接减半。

# 使用fp16量化
model_fp16 = ct.converters.onnx.convert(
    model="mobilenet_v2.onnx",
    minimum_ios_deployment_target='13',
    compute_precision=ct.precision.FLOAT16
)
model_fp16.save("MobileNetV2_fp16.mlmodel")

量化后的模型精度损失通常很小(1%以内),但推理速度能提升不少。我在一个图像分类项目里,量化后模型从30MB降到15MB,推理时间从50ms降到30ms。效果很明显。

总结一下:

PyTorch → ONNX → Core ML 这条路径,是目前最成熟、最稳定的方案。虽然中间多了一步,但每一步都有成熟的工具链支持,踩坑概率低。我个人强烈推荐你走这条路。

好了,模型转换的基础就讲到这里。你跟着代码走一遍,应该能顺利跑通。如果遇到问题,回头看看上面的表格,大部分坑我都帮你踩过了。

额外提示:转换完成后,记得用Xcode的Core ML模型预览功能检查一下。我每次都会看输入输出的数据类型和shape,确保和预期一致。这一步虽然简单,但能避免很多集成时的麻烦。

模型转换流程 PyTorch模型 .pth / .pt ONNX模型 .onnx Core ML模型 .mlmodel torch.onnx.export coremltools 关键参数与注意事项 • opset_version: 建议11或12 • dynamic_axes: 支持动态batch • minimum_ios_deployment_target: 建议iOS 13+ • 导出前必须调用 model.eval() • 用ONNX Runtime验证模型 • 量化可选 fp16 或 int8

可可精品资料库