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

♪ ♪

Pulkit：大家好 我是 Pulkit 一名 Core ML 团队的工程师 很高兴与大家分享 Core ML Tools 发布的多项更新 这些更新可以帮助你 优化机器学习模型的 大小和性能 随着模型能力的显著提升 越来越多的功能 开始由机器学习驱动 因此 每个 App 部署的 模型数量也随之增加 同时 App 中每个模型 也在不断增大 从而给 App 的大小 带来上升的压力 因此 控制模型大小十分重要 减小模型大小有以下几点好处 如果模型大小减小 你可以 在相同内存预算内传输更多模型 并且 传输模型的尺寸更大、 功能更强 同时 模型的运行速度也会更快 这是因为模型越小 内存和处理器之间传输的数据越少 所以 减小模型大小 似乎是个不错的选择 那么何种原因导致模型大小增大？ 我举一个例子来帮你理解一下 ResNet50 是一个时下流行的 图像分类模型 该模型第一层为卷积层 具有约 9000 个参数 而且 其共有 53 个 大小不一的卷积层 最后一层为线性层 具有约 210 万个参数 该模型总共有 多达 2500 万个参数 如果我使用 Float16 精度 来保存该模型 每个权重使用 2 个字节 那么模型大小为 50 兆字节 50 兆的模型很大 但当你接触 Stable Diffusion 等 更新的模型 你将得到更大的模型 现在 我们来谈谈减小 模型大小的方法 一种方法是设计 更为高效的模型架构 该架构可以用更少和更小的权重 实现良好的性能 另一种方法是压缩现有模型的权重 接下来 我会重点介绍 模型压缩这一方法 首先 让我们从介绍三种实用的 模型压缩技术开始 接着 我将演示两种 集成这些模型压缩技术的工作流 然后 我会讲解最新的 Core ML Tools 如何 帮助你将上述技术及工作流 应用到你的模型中 最后 Srijan 将会讨论 模型压缩对运行时性能的影响 首先 我们从压缩技术开始 有以下几种方法可以压缩模型权重 第一种方法是使用稀疏矩阵表示 对权重更高效地打包 此方法可通过剪枝技术实现 另一种方法是 降低用于存储权重的精度 此方法可通过量化及 调色板技术实现 但这两种方法均为有损压缩 并且相较于未经压缩的模型 利用此方法压缩的模型 通常精度也会更低

现在 我们逐一深入了解一下 这些技术

权重剪枝通过稀疏表示 可以帮助你 更加高效地打包模型权重 对权重矩阵进行稀疏或剪枝 就是将部分权重值设置为 0 这里是一个权重矩阵 要实现剪枝 我会将最小幅值权重设置为 0

现在 我只需要存储非零值 并且 我每引入一个零 就可以 节省大约 2 个字节的存储容量 当然了 我还需要存储零的位置 以便稍后重构稠密矩阵 随着稀疏程度不断增加 模型大小随之线性下降 稀疏度为 50% 的模型 意味着 50% 的权重都为零 对于 ResNet50 模型而言 其大小约为 28 兆字节 大约是 Float16 大小的一半

第二种权重压缩技术为量化 就是使用 8 位精度来存储权重 为了实现量化 你需要对权重值进行 缩放、移位和舍入 从而使其位于 INT8 范围内 在这个例子中 缩放比例为 2.35 从而最小值映射到 -127 偏差为 0 根据模型的不同 你还可以使用非零偏差 该方法有时可以帮助你 减小量化误差 缩放比例以及偏差 稍后将可用于将权重反量化 以使其恢复到原始范围 为了将权重精度降到 8 位以下 你可以使用权重聚类技术 也就是调色板技术 在该技术中 具有相似值的权重 会分到一组 并使用 所属聚类中心的值表示 这些聚类中心存储在一个查找表中 并且 原始权重矩阵 会转换为一个索引表 其中每个元素都会指向 相应的聚类中心 在该例子中 我共有四个聚类 所以 我可以只使用 2 个字节 来表示每个权重 从而 相较于 Float16 我可以实现 8 倍的压缩 权重可以使用 独特聚类中心的数量来表示 即 2 的 n 次方 其中 n 是调色板所使用的精度 所以 4 位调色板 表示你可以有 16 个聚类 虽然量化可以将模型大小减小一半 但调色板却可以 帮你实现高达 8 倍的压缩

总而言之 我们共有三种 用于权重压缩的技术 每种技术都使用了不同的方式 来表示权重 根据各自所使用的参数 这些技术可以实现不同程度的压缩 例如 剪枝使用的是稀疏度 而调色板使用的是字节数 现在 我将为你讲解如何将这些技术 整合到你的模型开发工作流中 首先 我们从 Core ML 模型转换的工作流开始 你可以使用喜欢的 python 训练框架训练一个模型 接着利用 Core ML Tools 将该模型转换为 Core ML 该工作流经过进一步扩展 可以成为训练后压缩的工作流 为此 你需要对训练后 经过转换的模型权重 增加一个压缩步骤 以减小模型整体大小 需要指出的是 该工作流可以从任一点开始 例如 你可以从预训练模型开始 而无需训练数据或是 从已经转换过的 Core ML 模型开始

在应用该工作流时 你可以对应用的压缩量进行选择 并且 你应用的压缩量越大 所得到的模型大小就越小 但如你所想 这其中存在一些权衡 具体来说 你刚开始使用的模型 未经压缩并具有一定精度 在对其进行压缩时 该模型大小会减小 同时精度也会受到影响 随着你提高压缩的程度 该影响会越来越明显 精度损失也会达到无法接受的程度

对于每个用例而言 该趋势及可接受的权衡各不相同 并且取决于具体的模型和数据集 为了向你展示该权衡 在实际应用中的情况 我们来看一个 用于分割图片对象的模型 在这张图片中 该模型会返回 每个像素属于沙发的概率 Float16 的基线模型 可以很好地分割对象 经过 10% 剪枝的模型 与基线模型的输出基本相似 稀疏度为 30% 时 瑕疵开始出现 并随着稀疏度的增加变得更加明显 而当我使用 40% 的剪枝时 模型完全崩溃 并且概率分布图也完全无法分辨 同样地 使用 8 位量化和 6 位调色板 可以保持基本模型的输出 在 4 位调色板中 你会看到一些瑕疵 而在 2 位调色板中 该模型则完全无法分割出对象 为了避免压缩率较高时 模型性能下降的情况 你可以使用另一个工作流 该工作流名为训练时压缩 在这里 你可以在压缩权重的同时 对模型中的部分数据进行微调 这种压缩采取渐进且可微分的方式 应用到模型中 以便权重可以重新适应 所施加的新约束 一旦模型达到了理想的精度 你便可以对其进行转换 获得一个经过压缩的 Core ML 模型 请注意 你既可以 在现有模型训练工作流中集成 训练时压缩 同时也可以从一个预训练模型开始 训练时压缩改善了 模型精度和压缩量之间的权衡 可以让你 在高压缩率的情况下 保持相同的模型性能 接下来 让我们再回到图像分割模型 对于训练时剪枝 模型输出可以 在稀疏度高达 40% 时保持不变 并且 该点也是后训练精度 开始下降的位置 事实上 即使稀疏度为 50% 和 75% 该模型输出的概率分布图 也与基线模型基本相似 当稀疏度达到 90% 时 你才可以观察到模型精度显著下降 同样地 使用训练时量化或调色板 模型输出可以在压缩率 高达 2 位时 仍然保持基线模型输出 回顾一下 你可以在模型转换 或模型训练时应用权重压缩 并且 后者可以以更长的训练时间 提供更好的精度权衡 由于第二个工作流 是在训练时进行压缩 所以我们需要使用可微分的操作 来实现压缩算法 现在 我们来探讨如何使用 Core ML Tools 来执行上述工作流 Core ML Tools 6 可以在 compression utils 子模块下 为后训练模型压缩 提供支持剪枝、调色板 及量化的 API 但是 在此之前 并没有针对训练时压缩的 API Core ML Tools 7 为训练时压缩 提供了具备相同功能的新 API 并且 我们将旧 API 和新 API 整合到同一个模块下 并将其命名为 coremltools.optimize 训练后压缩 API 迁移到 coremltools.optimize.coreml 下 而新的训练时 API 则位于 coremltools.optimize.torch 下 并且 后者可以与 PyTorch 模型 协同工作 首先 我们来详细看一下 训练后 API 在训练后压缩工作流中 输入为 Core ML 模型 该模型可通过 optimize.coreml 模块 中的三种方法进行更新 也就是应用我刚刚介绍过的 三种压缩技术 为了应用这些方法 首先 你需要创建一个 OptimizationConfig 对象 用来描述压缩模型时使用的方法 在这里 我使用幅值剪枝 并将目标稀疏度设置为 75% 定义完配置后 你便可以使用 prune_weights 方法 修剪模型 经过简单的一步式操作 便可以获得一个压缩模型 你也可以利用特定于上述技巧的配置 使用类似的 API 对权重应用调色板和量化 现在 让我们再来看看 训练时压缩的工作流 正如我之前介绍的 在这种情况下 你需要可训练模型以及数据 具体来说 想要使用 Core ML Tools 来压缩模型 首先 你需要一个 PyTorch 模型 最好带有预训练权重 接着 使用 optimize.torch 模块中 可用的 API 进行更新 便可以得到一个插入压缩层的 新 PyTorch 模型 然后 使用数据和 原始 PyTorch 训练代码对其进行微调 在这一步中 权重会得以调整 以用于接下来的压缩 并且 你可以在本地 MacBook 上 使用 MPS PyTorch 后端完成该步骤 一旦模型经过训练并恢复精度 你便可以将其 转换为 Core ML 模型 接下来 我们通过一段 代码示例更深入地进行介绍 首先 我从对想要压缩的模型 进行微调所需的 PyTorch 代码开始 你只需添加几行代码 便可轻松利用 Core ML Tools 添加训练时剪枝 第一步 你需要创建 MagnitudePrunerConfig 对象 描述你想要如何修剪模型 在这里 我将目标稀疏度设置为 75% 你也可以将配置写入 yaml 文件 并使用 from_yaml 方法进行加载 接着 使用想要压缩的模型 和刚刚创建的配置 创建一个剪枝对象 然后 调用 prepare 方法 在模型中插入剪枝层 在微调模型时 你还可以调用 step API 更新剪枝层的内部状态 在训练结束后 你可以调用 finalize 将修剪遮罩合并到权重中 接着使用转换 API 该模型便可以转换为 Core ML 该工作流同样适用于量化和调色盘 接下来 Srijan 将向你演示如何使用 Core ML Tools API 对物体检测模型应用调色盘 Srijan：谢谢你 Pulkit 我是 Srijan 接下来 我将向你演示 Core ML Tools 中的 optimize API 接下来 我会使用一个 SSD 模型 并将 ResNet18 作为骨干网络 来检测图片中的人像 首先 我们导入一些基础模型 和训练工具 接下来 我会使用刚刚谈到的 SSD ResNet18 模型举一个实例 为了简化操作 我将直接调用 预先编写好的 get_ssd_model 实用程序 现在 模型已加载完成 我们先来训练几轮 由于这是一个物体检测模型 因此该训练旨在 降低检测任务过程中 SSD 的损失 为了方便起见 train_epoch 实用程序 压缩了每轮训练模型所需的代码 比如通过不同的批次 调用 forward 函数、 计算损失值以及执行梯度下降 在训练过程中 SSD 损失逐渐下降 所以 现在我将该模型转换为 Core ML 模型 为此 我会先跟踪模型 然后调用 coremltools.convert API 接着 我们调用一个导入的实用程序 来检查模型大小

该模型大小为 23.6 兆字节 现在 我会在 Core ML 模型上 运行预测器 我选了一张在伦敦旅行时的照片 和一张其他照片来测试检测器 并且 我将模型中物体检测的 置信度阈值设为 30% 这样该模型只会对当前 置信度超过 30% 的 对象绘制边界框

检测效果似乎很准确 现在 我很好奇如果 减小该模型大小 结果又会如何 我首先来试一下训练后调色板 为此 我会从 coremltools.optimize.coreml 中 导入一些配置类和方法

接着 使用 6 位 对模型权重应用调色板 为了实现这点 我会创建一个 OpPalettizerConfig 对象 将模式设为 k 均值 并将 nbits 设为 6 并且该设置指定操作级别的参数 从而我可以对每个操作 应用不同的调色板 但现在 我会对所有操作应用相同的 6 位调色板模式 定义 OptimizationConfig 并将 op_config 作为全局参数 传递给该对象就可以了

接着 优化配置就会和 转换后的模型一同传递给 palettize_weights 以获得调色板模型 现在 我们再来看看 模型大小缩小至多少

可以看到 模型大小 减小了约 9 兆字节 但这会影响测试图片的性能吗？ 我们一起来看看 哇 检测仍然有效 十分激动 希望接下来使用 2 位 训练后调色盘时 也这么幸运 这一步非常简单 在 OpPalettizerConfig 中 将 nbits 从 6 改为 2 接着再次运行 palettize_weights API

我们使用实用程序看看当前 Core ML 模型的大小和性能

如我们所料 模型大小已减小至约 3 兆字节 但是 性能好像不太理想 它在两张图片中都无法检测到人像 预测器中没有出现边界框 这就表示在模型所预测的 边界框中 没有一个置信度 高于 30% 的阈值 我们再来试试 2 位训练时调色盘 看看其性能是否会更加出色

首先 我会从 coremltools.optimize.torch 中 导入 DKMPalettizerConfig 以及 DKMPalettizer DKM 是一种学习 权重聚类的算法 该算法会对权重进行 基于注意力的可微分的 k 均值操作 现在 我们来定义调色板配置 只需要在 global_config 中 将 n_bits 设为 2 接着 所有受支持的模块 都会使用 2 位调色板 然后在这里 我将从模型和配置中 创建一个调色板对象 调用 prepare API 便可以在模型中插入 调色板友好的模块 现在 我们来对模型微调几轮 现在 模型已经完成了微调 我来调用 finalize API 将调色板权重恢复为模型的权重 这样整个过程便完成了 下一步 我们来检查模型大小 为此 我需要将 torch 模型 转换为 Core ML 模型 我们首先使用 torch.jit.trace 对模型进行跟踪 这一次 我不仅会调用 convert API 还会使用 PassPipelin 额外标志 并将其值设置为 DEFAULT_PALETTIZATION 这表明转换器会对转换权重 使用调色板的表示

让我们来看看模型大小 以及在测试图片上的性能 可以看到 使用训练时调色板后 模型大小约为 3 兆字节 相当于 8 倍压缩 但不同于训练后调色板模型 该模型可以在测试图片上 实现正确检测 由于这只是一个演示 我只使用 两张示例图片来测试模型性能 在实际情况下 我会使用类似 均值平均精度这样的矩阵 在验证数据集上进行评估

让我们回顾一下 首先 我从一个训练过的模型开始 经过转换 得到一个 大小为 23.6 兆字节 并带有 Float16 权重的模型 接着 我利用 palettize_weights API 快速得到了一个尺寸更小且 带有 6 位权重的模型 并且该模型在运行数据时表现良好 但是 我将该模型 继续压缩到 2 位时 其性能显著下降 之后 我又利用 optimize.torch API 对 torch 模型进行了更新 并利用可微分 k 均值算法 对其进行了几轮微调 这样 在 2 位压缩选项下 我便可以得到良好的精度 尽管该演示使用了特定模型和 优化算法的组合 但是该工作流可以 推广到你的用例中 并帮助你清楚权衡 期望的压缩量 和重新训练模型所需要的 时间和数据 我们来谈谈最后一个话题 性能 在这里 我会简单介绍 最近对 Core ML 运行时 进行的改进 以便你在 App 中部署时 更高效地执行此类模型 首先 我们来看看 在 iOS 16 和 iOS 17 中 运行时的几点关键差异 iOS 16 支持仅对权重 进行压缩的模型 而在 iOS 17 中 8 位激活量化模型 也可以得到执行 在 iOS 16 中 权重压缩模型的运行速度 与相应带有 float 权重的模型相同； 而在 iOS 17 中 Core ML 运行时经过更新后 压缩模型在某些情况下 可以运行得更快 较新版本的 macOS、Apple tvOS 和 watchOS 现也支持类似的运行时改进 但这些改进是如何实现的呢？ 对于仅压缩权重的模型而言 由于激活值是浮点精度 因此 在进行类似卷积 或矩阵乘法操作前 需要对权重值进行解压 以匹配其他输入的精度 在 iOS 16 运行时中 解压的步骤会提前进行 因此 在这种情况下 模型便会转换为 完全的 float 精度模型 并存储在内存中以待执行 从而 推理反应时间不会发生改变 但在 iOS 17 中 在某些特定情况下 权重会恰好在操作执行前进行解压 这样就可以从内存中 加载位数较小的权重 但是需要在 每次调用推理时进行解压 对于神经网络引擎 以及某些内存受限的模型 等运算单元来说 这么做还可以带来推理性能的提升 为了详细说明这些运行时的优势 我选择并概述了几个模型 并与 Float16 变体进行对比 绘制了推理速度的相对增加量 正如我们所预期的那样 速度增加量取决于模型和硬件 这些是 iPhone 14 Pro Max 上 4 位调色板模型的速度增加范围 改进幅度均在 5% 到 30% 之间 对于稀疏模型而言 其改进幅度也取决于模型类型 并且某些模型的运行速度 比同类型的 Float16 变体快了 75% 那么问题来了： 获得最佳反应时间性能 的策略是什么？ 最好的方法就是 从 float 模型开始 然后利用 optimize.coreml API 探索模型的各种表示方式 这个过程非常快速 因为无需重新训练模型 接着 在你感兴趣的设备上 对其进行分析 针对这点 Xcode 中的 Core ML 性能报告便可以为你提供 有关推理的大量信息 包括操作运行的位置等 然后 根据能够为你带来 最佳性能的配置进行筛选 在这之后 你便可以专注于评估精度 并尝试改进 但在你确定最终模型前 你需要使用 torch 和 Core ML Tools 应用一些训练时压缩

总而言之 减小模型大小至关重要 并且 借助全新 Core ML Tools API 你可以更加轻松地实现该目标 同时还能实现更低的内存占用 以及更快速的推理 欲查看更多选项和基准数据 请访问我们的文档 同时 我还强烈建议你 观看“通过异步预测改进 Core ML 集成”讲座 其中介绍了我今天 幻灯片中没有涵盖的 Core ML 框架的改进 谢谢大家 祝你拥有愉快的压缩体验