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

♪ ♪ Pulkit：大家好，我是 Pulkit，是 Core ML 团队的工程师。我很高兴与大家分享 Core ML Tools 的一些更新。这些更新可帮助您优化机器学习模型的大小和性能。随着模型能力的显着提高，越来越多的功能正在由机器学习驱动。因此，单个应用程序中部署的模型数量不断增加。与此同时，应用程序中的每个模型也变得越来越大，这给应用程序的大小带来了向上的压力。因此，控制模型大小至关重要。减小模型尺寸有几个好处。如果每个模型更小，您可以在相同的内存预算中运送更多模型。它还可以帮助您交付更大、功能更强大的模型。它还可以帮助使模型运行得更快。这是因为较小的型号意味着在内存和处理器之间移动的数据较少。因此，减小模型的尺寸似乎是个好主意。是什么让模型变大？让我通过一个例子来帮助您理解。 ResNet50是一种流行的图像分类模型。它的第一层是一个大约有 9000 个参数的卷积层。它总共有 53 个大小不同的卷积层。最后，它有一个包含约 210 万个参数的线性层。所有这些加起来总共有 2500 万个参数。如果我使用 Float16 精度保存模型，则每个权重使用 2 个字节，并且我得到一个大小为 50 MB 的模型。 50 MB 的模型很大，但是当您使用稳定扩散等较新的模型时，您最终会得到更大的模型。现在，我们来谈谈获得更小模型的一些途径。一种方法是设计一种更高效的模型架构，可以用越来越小的权重实现良好的性能。另一种方法是压缩现有模型的权重。模型压缩的这条路径是我将重点关注的。我将首先描述三种有用的模型压缩技术。接下来，我将演示集成这些模型压缩技术的两个工作流程。然后，我将说明最新的 Core ML 工具如何帮助您将这些技术和工作流程应用到您的模型中。最后，Srijan 将讨论模型压缩对运行时性能的影响。让我们从压缩技术开始。有几种压缩模型权重的方法。第一种方法是使用稀疏矩阵表示更有效地打包它们。这可以通过使用称为修剪的技术来实现。另一种方法是降低用于存储权重的精度。这可以通过量化或调色来实现。这两种策略都是有损的，并且与未压缩的模型相比，压缩模型的准确度通常稍差。

现在让我们更深入地了解这些技术。

权重修剪可帮助您使用稀疏表示有效地打包模型权重。稀疏或修剪权重矩阵意味着将一些权重值设置为 0。我从权重矩阵开始。为了修剪它，我可以将最小量值权重设置为 0。

现在，我只需要存储非零值。每引入一个零，我最终都会节省大约 2 个字节的存储空间。当然，我还需要存储零的位置，以便稍后重建稠密矩阵。模型大小随着引入的稀疏度水平线性下降。 50% 稀疏模型意味着其 50% 的权重为零，对于 ResNet50 模型，其大小约为 28 MB，大约是 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 中提供了训练后模型压缩 API，用于压缩 utils 子模块下的修剪、调色和量化。然而，没有用于训练时间压缩的 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 训练代码对其进行微调。这是调整权重以允许压缩的步骤。您可以使用 MPS PyTorch 后端在本地 MacBook 上执行此步骤。一旦模型经过训练以重新获得准确性，请将其转换为 Core ML 模型。让我们通过代码示例进一步探讨这一点。我从微调我想要压缩的模型所需的 PyTorch 代码开始。您只需添加几行代码即可轻松利用 Core ML Tools 来添加训练时间修剪。您首先创建一个 MagnitudePrunerConfig 对象来描述如何修剪模型。在这里，我将目标稀疏度设置为 75%。您还可以将配置写入 yaml 文件并使用 from_yaml 方法加载它。然后，您使用要压缩的模型和刚刚创建的配置创建一个修剪器对象。接下来，您调用prepare在模型中插入修剪层。在微调模型时，您可以调用步骤 API 来更新剪枝器的内部状态。训练结束时，您调用 Finalize 将修剪蒙版折叠到权重中。然后可以使用转换 API 将该模型转换为 Core ML。相同的工作流程也可用于量化和调色。现在，Srijan 将引导您完成一个演示，展示如何使用 Core ML Tools API 来调色板对象检测模型。 Srijan：谢谢你，普尔基特。我叫 Srijan，我将带您演示 Core ML Tools 优化 API。我将使用带有 ResNet18 主干的 SSD 模型来检测图像中的人物。让我们首先导入一些基本模型和训练实用程序。我将首先获取我刚才谈到的 SSD ResNet18 模型的实例。为了简化事情，我将为此调用预先编写的 get_ssd_model 实用程序。现在模型已加载，让我们对其进行几个时期的训练。由于它是一个目标检测模型，因此训练的目标是减少检测任务的 SSD 损失。为了简洁起见，train_epoch 实用程序封装了训练一个 epoch 模型所需的代码，例如通过不同批次调用前向、计算损失和执行梯度下降。在训练过程中，SSD 损失似乎有所下降。我现在将该模型转换为 Core ML 模型。为此，我将首先跟踪模型，然后调用 coremltools.convert API。让我们调用一个导入的实用程序来检查模型的大小。

模型大小为 23.6 MB。现在，我将在 Core ML 模型上运行预测。我选择了一张我自己在伦敦旅行中的图像以及另一张图像来测试检测。模型检测对象的置信阈值设置为 30%，因此它只会绘制对存在对象至少有 30% 置信度的框。

这个检测似乎很准确。我现在很好奇是否可以减小这个模型的尺寸。我将首先尝试训练后调色。为此，我将从 coremltools.optimize.coreml 导入一些配置类和方法。

我现在要用 6 位对模型的权重进行调色。为此，我将创建一个 OpPalettizerConfig 对象，将模式指定为 kmeans，将 nbits 指定为 6。这将在操作级别指定参数，并且我可以对每个操作进行不同的调色。然而，现在，我将对所有操作应用相同的 6 位模式。我将通过定义 OptimizationConfig 并将此 op_config 作为全局参数传递给它来实现这一点。

然后，优化配置与转换后的模型一起传递给 papetize_weights 方法，以获得调色模型。让我们看看现在尺寸缩小到什么程度了。

模型的大小已降至 9 MB 左右，但这是否影响了测试图像的性能？让我们来看看吧。哇，检测仍然有效。我真的很高兴现在能尝试 2 位训练后调色。执行此操作非常简单，只需在 OpPalettizerConfig 中将 nbits 从 6 更改为 2 并再次运行 Palettize_weights API 即可。

让我们使用实用程序来查看此 Core ML 模型的大小和性能。

正如预期的那样，模型的大小已减小至 3 MB 左右。然而，性能并不理想，因为模型无法检测到两张图像中的人物。预测中没有出现框，因为模型预测的框的置信概率均不高于 30% 阈值。让我们尝试一下 2 位训练时间调色，看看效果是否更好。

我将首先从 coremltools.optimize.torch 导入 DKMPalettizerConfig 和 DKMPalettizer 来执行此操作。 DKM 是一种通过对权重簇执行基于注意力的可微分 kmeans 操作来学习权重簇的算法。现在是时候定义调色板配置了。只需在 global_config 中将 n_bits 指定为 2，所有支持的模块都将被 2 位调色。在这里，我将根据模型和配置创建一个调色机对象。现在调用准备 API 会将调色板友好的模块插入模型中。是时候对模型进行几个时期的微调了。现在模型已经微调，我将调用 Finalize API，它将恢复调色板权重作为模型的权重，从而完成该过程。下一步是检查模型的尺寸。为此，我将把 torch 模型转换为 Core ML 模型。让我们首先使用 torch.jit.trace 跟踪模型。我现在将调用转换 API，这一次，我将使用一个名为 PassPipeline 的附加标志并将其值设置为 DEFAULT_PALETTIZATION。这将指示转换器使用转换后的权重的调色板表示。

让我们看看模型的大小及其在测试图像上的性能。我可以看到，训练时调色的模型也约为 3 MB，压缩率为 8 倍，但与训练后调色的模型不同，该模型正确地对测试图像执行检测。由于这是一个演示，我只是在两个示例图像上测试了模型性能。在现实场景中，我会使用平均精度等指标并在验证数据集上进行评估。

让我们回顾一下。我从一个经过训练的模型开始，并将其转换为具有 Float16 权重的 23.6 MB 模型。然后，我使用 papetize_weights API 快速获得具有 6 位权重的较小模型，该模型在我的数据上表现良好。然而，当我进一步将其推至 2 位时，性能明显下降。在此之后，我使用 optimize.torch API 更新了 torch 模型，并使用可微分 kmeans 算法对几个时期进行了微调。这样，我就能够通过 2 位压缩选项获得良好的精度。虽然演示采用了特定的模型和优化算法组合，但此工作流程将推广到您的用例，并将帮助您找到所需的压缩量与重新训练模型所需的时间和数据之间的权衡。这引出了我们的最后一个主题：性能。我想简要介绍一下对 Core ML 运行时所做的改进，以便在部署到应用程序中时更有效地执行此类模型。让我们看看 iOS 16 和 iOS 17 中运行时之间的一些关键区别。在 iOS 16 中，支持仅权重压缩模型，而在 iOS 17 中，还可以执行 8 位激活量化模型。在 iOS 16 中，权重压缩模型的运行速度与具有浮动权重的相应模型相同，而在 iOS 17 中，Core ML 运行时已更新，现在压缩模型在某些场景下运行得更快。新版本的 macOS、tvOS 和 watchOS 中也提供了类似的运行时改进。但这些改进是如何实现的呢？在仅压缩权重的模型中，由于激活采用浮点精度，因此在发生卷积或矩阵乘法等操作之前，需要解压缩权重值以匹配其他输入的精度。此解压步骤在 iOS 16 运行时中提前进行。因此，在这种情况下，模型在执行之前会在内存中转换为完全浮点精度模型。因此，推理延迟没有变化。然而，在 iOS 17 中，在某些情况下，权重会在执行操作之前及时解压缩。这样做的优点是可以从内存中加载较小的位权重，但代价是在每次推理调用中进行解压缩。对于某些计算单元（例如神经引擎）和某些受内存限制的模型类型，这可能会带来推理增益。为了说明这些运行时优势，我选择并分析了几个模型，并绘制了与 Float16 变体相比，它们的推理速度加快的相对量。

正如预期的那样，加速量取决于模型和硬件。这些是 iPhone 14 Pro Max 上 4 位调色板模型的加速范围。改进幅度在 5% 到 30% 之间。对于稀疏模型，根据模型类型也有不同的改进，一些模型的运行速度比 Float16 变体快 75%。现在的问题是：获得最佳延迟性能的策略是什么？这将从浮点模型开始，并使用 optimization.coreml API 来探索模型的各种表示形式。这会很快，因为它不需要重新训练模型。然后，在您感兴趣的设备上对其进行配置。为此，Xcode 中的 Core ML 性能报告将为您提供大量推理可见性，包括操作运行的位置。然后，根据哪些配置为您带来最佳收益，列出候选名单。之后，您可以专注于评估准确性并尝试改进，这可能需要在最终确定模型之前使用 torch 和 Core ML Tools 进行一些训练时间压缩。