Use Core ML Tools for machine learning model compression
Machine Learning & AI 进阶 20m

使用 Core ML Tools 压缩机器学习模型

Use Core ML Tools for machine learning model compression

2023年6月5日

在 Apple 官方观看视频

一句话判断

如果你的 App 里塞了多个 Core ML 模型导致体积膨胀，这场 Session 提供了三种压缩手段——调色板量化（Palette Quantization）、线性量化（Linear Quantization）和结构化剪枝（Structured Pruning），配合优化工作流实现精度和体积的平衡。

这场 Session 讲了什么

Session 围绕 Core ML Tools 的模型压缩能力展开，介绍了三种压缩技术和两个优化工作流。

模型压缩的动机很直接：模型越来越大（ResNet50 的 50MB 已经算小的，Stable Diffusion 更是巨大），App 中部署的模型数量也在增加。压缩模型不仅能减少 App 体积，还能减少内存与处理器之间的数据搬运，间接提升推理速度。

调色板量化（Palette Quantization）：将浮点权重聚类到有限的离散值集合中（类似图像的颜色量化）。比如把 65536 个 Float16 值映射到 256 个聚类中心，索引只需 1 字节存储，压缩比约 4 倍。聚类数量可以选择 4、8、16…256，灵活性很高。

线性量化（Linear Quantization）：将浮点权重线性映射到整数范围（如 Int8），每个权重只需 1 字节，压缩比稳定在 2 倍。相比调色板方式更简单直接，适合大模型。

结构化剪枝（Structured Pruning）：直接移除对输出影响最小的权重通道或层，最高可实现 2-4 倍压缩。被剪枝的结构在推理时完全不参与计算，所以还能提升速度。

两个优化工作流：训练后量化（Post-Training Quantization）直接压缩已有模型，简单快速；训练时量化感知（Quantization-Aware Training）在训练过程中考虑量化误差，精度损失更小。

值得深挖的点

压缩方式的选择：调色板量化适合中小模型、对精度敏感的场景（可以精细控制聚类数量）；线性量化适合大模型（操作简单，压缩比稳定）；剪枝适合同时需要减少体积和推理时间的场景。

精度损失的权衡：压缩必然带来精度损失，但实际影响取决于模型和任务。Session 提供的优化工作流允许你在压缩前后对比精度，找到最佳平衡点。对于分类任务，1-2% 的精度下降通常可以接受。

剪枝的加速效果：量化和调色板方法只减少存储体积，推理时仍需反量化计算。剪枝直接移除计算结构，在推理速度上的收益更直接。

代码片段

使用 Core ML Tools 进行模型压缩：

import coremltools as ct
from coremltools.optimize.coreml import (
    OpLinearQuantizerConfig,
    LinearQuantizer,
    OpPaletteQuantizerConfig,
    PaletteQuantizer
)

# 加载原始模型
model = ct.models.MLModel("resnet50.mlmodel")

# 方法一：线性量化到 Int8
linear_config = OpLinearQuantizerConfig(
    mode="linear_symmetric",  # 对称量化
    weight_threshold=2048      # 只量化大于此阈值的权重
)
quantizer = LinearQuantizer(global_config=linear_config)
quantized_model = quantizer.compress(model)

# 方法二：调色板量化
palette_config = OpPaletteQuantizerConfig(
    nbits=8,  # 256 个聚类中心，每个索引 1 字节
    mode="kmeans"
)
palette_quantizer = PaletteQuantizer(global_config=palette_config)
compressed_model = palette_quantizer.compress(model)

# 保存压缩后的模型
compressed_model.save("resnet50_compressed.mlmodel")

最佳实践

  • 先用训练后量化快速验证压缩效果，如果精度损失可接受就直接上线
  • 对精度要求高的模型使用量化感知训练
  • 调色板量化从 8-bit 开始尝试，逐步降低到 4-bit 观察精度变化
  • 结构化剪枝适合部署在移动端的模型——同时减少体积和推理时间
  • 压缩后务必在真实数据集上评估精度，不要只看压缩比

还有什么值得关注

  • 量化感知训练的具体实现方式和集成步骤
  • 不同压缩技术在 Neural Engine 和 GPU 上的推理性能差异
  • Stable Diffusion 等生成式模型的压缩策略
  • Core ML Tools 的分析工具——如何在压缩前后可视化权重分布
WWDC 2023