更快更稳定的Core ML在iPhone、iPad和Mac上的扩散

WWDC’23（苹果全球开发者大会）上周举行。很多新闻都集中在主题演讲中的Vision Pro发布上，但其中还有更多内容。与往年一样，WWDC周提供了超过200个技术会议，深入介绍了苹果操作系统和框架即将推出的特性。今年，我们特别关注Core ML中关于压缩和优化技术的改进。这些改进使得运行Stable Diffusion等模型更快、内存使用更少！作为一个例子，考虑以下我去年12月在iPhone 13上进行的测试，与目前使用6位调色板化速度相比：

更快更稳定的Core ML在iPhone、iPad和Mac上的扩散四海第1张 iPhone上的Stable Diffusion，去年12月和现在使用6位调色板化

新的Core ML优化

Core ML是一个成熟的框架，可以在设备上高效地运行机器学习模型，充分利用苹果设备的计算硬件：CPU、GPU和专门用于机器学习任务的神经引擎。设备上的执行正经历着一个非常有趣的时期，这是由于Stable Diffusion和带有聊天界面的大型语言模型的流行所引起的。许多人出于各种原因（包括便利性、隐私和API成本节约）希望在自己的硬件上运行这些模型。自然而然地，许多开发人员正在探索在设备上高效运行这些模型的方法，并创建新的应用和用例。对于社区来说，对实现这一目标做出贡献的Core ML改进是一个重大的消息！

Core ML优化改进包括两个不同（但相辅相成）的软件包：

Core ML框架本身。这是在苹果硬件上运行ML模型的引擎，是操作系统的一部分。模型必须以框架支持的特殊格式导出，这个格式也被称为“Core ML”。
coremltools转换包。这是一个开源的Python模块，其任务是将PyTorch或Tensorflow模型转换为Core ML格式。

coremltools现在包含一个名为coremltools.optimize的新子模块，其中包含所有压缩和优化工具。有关此软件包的详细信息，请参阅此WWDC会议。在Stable Diffusion的情况下，我们将使用6位调色板化，这是一种将模型权重从16位浮点表示压缩为每个参数仅使用6位的量化类型。名称“调色板化”是指类似于计算机图形学中使用的一种技术，该技术用于使用有限的一组颜色工作：颜色表（或“调色板”）包含固定数量的颜色，并且图像中的颜色被替换为调色板中可用的最接近的颜色的索引。这立即提供了大大减小存储大小的好处，从而减少下载时间和设备上的磁盘使用。

更快更稳定的Core ML在iPhone、iPad和Mac上的扩散四海第2张 2位调色板化的说明。图片来源：苹果WWDC’23会议使用Core ML工具进行机器学习模型压缩。

压缩的6位权重不能用于计算，因为它们只是一个表中的索引，不再代表原始权重的大小。因此，Core ML在使用之前需要解压缩调色板化的权重。在之前的Core ML版本中，解压缩是在模型首次从磁盘加载时进行的，所以使用的内存量等于未压缩模型的大小。通过新的改进，权重被保留为6位数，并在推理从层到层进行时动态转换。这可能看起来很慢——一个推理运行需要大量的解压缩操作——但它通常比准备所有权重以16位模式更高效！原因是内存传输是执行的关键路径，传输较少的内存比传输未经压缩的数据更快。

使用量化和优化的Stable Diffusion模型

去年12月，苹果推出了ml-stable-diffusion，这是一个基于扩散器的开源存储库，可以轻松将Stable Diffusion模型转换为Core ML。它还对变压器的注意力层进行了优化，以在神经引擎上进行更快的推理（在支持该引擎的设备上）。在WWDC之后，ml-stable-diffusion已经更新，具体如下：

使用--quantize-nbits在转换过程中支持量化。您可以量化为8、6、4甚至2位！为了获得最佳结果，我们推荐使用6位量化，因为精度损失小，同时实现快速推理和显著的内存节省。如果您想要更低的位数，请查看此部分以获取高级技术。
对注意力层进行额外优化，以在神经引擎上获得更好的性能！诀窍是将查询序列分成512个片段，以避免创建大型中间张量。该方法在代码中称为SPLIT_EINSUM_V2，可以提高性能10%至30%。

为了使每个人都能轻松利用这些改进，我们已将四个官方稳定扩散模型转换并推送到Hub。这些都是变体：

请注意，每个变体都以Core ML格式和zip档案格式提供。Zip文件非常适合原生应用程序，例如我们的开源演示应用程序和其他第三方工具。如果您只想在自己的硬件上运行模型，最简单的方法是使用我们的演示应用程序并选择要测试的量化模型。您需要使用Xcode编译应用程序，但很快将提供下载更新的App Store。有关更多详情，请查看我们之前的帖子。

更快更稳定的Core ML在iPhone、iPad和Mac上的扩散四海第3张在演示应用程序中运行6位稳定扩散-2-1-base模型

如果您想要下载特定的Core ML包以集成到自己的Xcode项目中，您可以克隆存储库，或者只需使用以下代码下载所需的版本。

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-2-1-base-palettized"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"模型已下载至{model_path}")

转换和优化自定义模型

如果您想使用个性化的稳定扩散模型（例如，如果您已经微调或进行了自定义模型的训练），可以使用Apple的ml-stable-diffusion存储库进行转换。以下是您可以进行转换的简要概述，但我们建议您阅读文档的详细内容。

选择要转换的模型。您可以训练自己的模型或从Hugging Face Diffusers模型库中选择一个。例如，让我们转换prompthero/openjourney-v4。
安装apple/ml-stable-diffusion并使用ORIGINAL注意力实现运行第一个转换，如下所示：

python -m python_coreml_stable_diffusion.torch2coreml \
    --model-version prompthero/openjourney-v4 \
    --convert-unet \
    --convert-text-encoder \
    --convert-vae-decoder \
    --convert-vae-encoder \
    --convert-safety-checker \
    --quantize-nbits 6 \
    --attention-implementation ORIGINAL \
    --compute-unit CPU_AND_GPU \
    --bundle-resources-for-swift-cli \
    --check-output-correctness \
    -o models/original/openjourney-6-bit

重复对SPLIT_EINSUM_V2注意力实现进行转换：

python -m python_coreml_stable_diffusion.torch2coreml \
    --model-version prompthero/openjourney-v4 \
    --convert-unet \
    --convert-text-encoder \
    --convert-vae-decoder \
    --convert-safety-checker \
    --quantize-nbits 6 \
    --attention-implementation SPLIT_EINSUM_V2 \
    --compute-unit ALL \
    --bundle-resources-for-swift-cli \
    --check-output-correctness \
    -o models/split_einsum_v2/openjourney-6-bit

在所需硬件上测试转换后的模型。通常情况下，在macOS上ORIGINAL版本的效果更好，而在iOS上SPLIT_EINSUM_V2通常更快。有关更多详细信息和额外数据点，请参阅社区在以前版本的 Stable Diffusion for Core ML 上贡献的这些测试。
要将所需的模型集成到您自己的应用程序中：
- 如果您要在应用程序内分发模型，请使用.mlpackage文件。请注意，这将增加您的应用程序二进制文件的大小。
- 否则，您可以使用编译的Resources在应用程序启动时动态下载它们。

使用少于6位

6位量化是模型质量、模型大小和便利性之间的最佳平衡点 – 您只需提供一个转换选项，就能够对任何预训练模型进行量化。这是一个后训练压缩的示例。

上周发布的coremltools的测试版本还包括训练时压缩方法。这里的想法是，您可以在微调预训练的 Stable Diffusion 模型并在微调过程中执行权重压缩。这样，您可以使用4位甚至2位压缩，同时最小化质量损失。之所以能够做到这一点，是因为使用可微分的算法执行了权重聚类，因此我们可以应用常规的训练优化器来查找量化表，同时最小化模型损失。

我们计划很快评估这种方法，并迫不及待地想看看4位优化模型的工作效果和运行速度。如果您比我们更早尝试了，请给我们留言，我们将很乐意核实一下 🙂

结论

量化方法可用于减小 Stable Diffusion 模型的大小，在设备上运行更快并消耗更少资源。Core ML和coremltools的最新版本支持6位调色板等易于应用且对质量影响最小的技术。我们已将6位调色板模型添加到Hub中，这些模型足够小，可以在iOS和macOS上运行。我们还展示了如何自己转换微调模型，并迫不及待地想看看您如何使用这些工具和技术！

更快更稳定的Core ML在iPhone、iPad和Mac上的扩散

目录

新的Core ML优化

使用量化和优化的Stable Diffusion模型

转换和优化自定义模型

使用少于6位

结论