4、Core ML模型封装:使用coremltools进行模型转换与配置、输入输出设置

好,咱们继续。前面我们把MobileViT模型训练好了,也做了量化。但模型还是PyTorch的格式,iOS不认识。所以这一步,我们要把它变成Core ML的格式。

说白了,就是给iOS写个翻译官。让iPhone能读懂我们的模型。

4.1 为什么非要转成Core ML?

你可能想问:我直接用PyTorch的模型不行吗?嗯,理论上可以,但实际体验会很糟糕。

Core ML是苹果自家的推理引擎。它有几个好处:

  • 硬件加速:自动调用ANE(苹果神经网络引擎)
  • 内存优化:系统级的内存管理,不容易爆
  • 低功耗:比纯CPU推理省电得多
  • 无缝集成:Xcode里直接拖进去就能用

我在项目中遇到过,有人直接用LibTorch跑模型,结果iPhone发烫严重,电池哗哗掉。换成Core ML后,温度直接降了10度。所以这一步,别偷懒。

4.2 coremltools是什么?

coremltools是苹果官方出的转换工具。它能把PyTorch、TensorFlow、ONNX的模型,转成Core ML格式(.mlmodel或.mlpackage)。

安装很简单:

pip install coremltools

我个人习惯用最新版。因为苹果每年都在更新Core ML的规范,旧版本可能不支持新特性。

4.3 模型转换的基本流程

转换的核心就三步:加载模型 → 设置输入输出 → 保存。

先看一个最简单的例子:

import coremltools as ct
import torch

# 1. 加载PyTorch模型
model = MobileViT(num_classes=10)
model.load_state_dict(torch.load('mobilevit.pth'))
model.eval()

# 2. 准备一个示例输入
example_input = torch.randn(1, 3, 256, 256)

# 3. 转换为Core ML
traced_model = torch.jit.trace(model, example_input)
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name='input_image', shape=example_input.shape)]
)

# 4. 保存
mlmodel.save('MobileViT.mlpackage')

你看,核心就一个ct.convert()调用。但这里有几个坑,我得跟你说说。

⚠️ 注意: 转换时一定要用torch.jit.trace,而不是torch.jit.script。因为MobileViT里有些动态控制流,script会报错。trace更稳定。

4.4 输入输出的精细配置

刚才的例子是最简配置。但实际项目中,我们需要更精细的控制。

4.4.1 输入类型选择

Core ML支持两种输入类型:

类型 适用场景 特点
ImageType 图像分类、目标检测 自动处理颜色空间、缩放
TensorType 特征向量、序列数据 原始数值,不做预处理

我个人建议:如果输入是图像,就用ImageType。因为它会自动帮你做归一化,省得在iOS端再写一遍。

# 图像输入 - 推荐
inputs=[ct.ImageType(
    name='input_image',
    shape=(1, 3, 256, 256),
    scale=1.0/255.0,  # 归一化
    bias=[0.0, 0.0, 0.0]
)]

# 张量输入 - 如果你要自己处理
inputs=[ct.TensorType(
    name='input_tensor',
    shape=(1, 3, 256, 256)
)]

4.4.2 输出设置

输出通常有两种情况:

  • 分类结果:用ClassifierConfig
  • 特征向量:直接输出张量
# 分类输出
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name='input_image', shape=(1, 3, 256, 256))],
    classifier_config=ct.ClassifierConfig(
        class_labels=['cat', 'dog', 'bird']
    )
)

# 特征输出(比如用于特征提取)
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name='input_image', shape=(1, 3, 256, 256))],
    outputs=[ct.TensorType(name='features')]
)
💡 小技巧: 如果你不确定输出是什么,可以先不设置outputs,让coremltools自动推断。转换完成后,用mlmodel.get_spec()查看输入输出信息。

4.5 我踩过的坑

嗯,这里分享几个我实际项目中遇到的坑,你注意避开。

坑1:动态尺寸问题

MobileViT支持任意尺寸输入,但Core ML要求固定尺寸。我曾经没固定尺寸,结果转换出来的模型在iPhone上跑不起来。

解决办法:要么固定输入尺寸,要么用ct.utils.make_pipeline做尺寸适配。

坑2:颜色通道顺序

PyTorch默认是RGB,但Core ML的ImageType默认也是RGB。按理说没问题。但我有一次用了OpenCV读图,它是BGR的,结果颜色全乱了。

解决办法:在iOS端用UIImage或CVPixelBuffer时,确保颜色空间是RGB。

坑3:模型太大

MobileViT虽然轻量,但转换后的.mlpackage可能比原始模型大。我试过,一个20MB的PyTorch模型,转成Core ML后变成了60MB。

解决办法:用ct.models.neural_network.quantization_utils做量化,能压到20MB以内。

4.6 验证转换结果

转换完别急着用。先验证一下:

# 加载转换后的模型
mlmodel = ct.models.MLModel('MobileViT.mlpackage')

# 准备测试数据
import numpy as np
test_image = np.random.randn(1, 3, 256, 256).astype(np.float32)

# 预测
prediction = mlmodel.predict({'input_image': test_image})
print(prediction)

如果输出和PyTorch模型一致,那就稳了。如果不一致,检查一下预处理参数。

📌 核心要点:
  • 用torch.jit.trace而不是script
  • 图像输入用ImageType,省去iOS端预处理
  • 固定输入尺寸,避免动态尺寸问题
  • 转换后一定要验证输出一致性

好了,模型转换这块就这些。说白了就是三步:trace、convert、save。但细节决定成败,尤其是输入输出的配置,一定要和iOS端的预处理对齐。

下一节我们讲怎么在Xcode里集成这个模型,以及如何用Swift调用它做推理。到时候你会发现,前面这些准备工作做得越扎实,后面就越省心。

Core ML模型封装流程 步骤1 加载PyTorch模型 步骤2 torch.jit.trace 步骤3 ct.convert() 保存 配置细节 输入配置 • ImageType:自动处理颜色空间 • TensorType:原始数值输入 • 固定尺寸:避免动态尺寸问题 输出配置 • ClassifierConfig:分类结果 • TensorType:特征向量 • 自动推断:不设置outputs 验证环节 • 用mlmodel.predict()测试输出一致性 • 对比PyTorch和Core ML的输出,确保误差在可接受范围内 核心:trace → convert → save,细节决定成败

可可精品资料库