优化稳定扩散：使用NNCF和🤗 Optimum对Intel CPU进行优化
优化稳定扩散：使用NNCF和🤗 Optimum优化Intel CPU
潜在扩散模型在解决文本到图像生成问题时具有革命性的作用。 稳定扩散是社区和行业中广泛采用的最著名的例子之一。稳定扩散模型的思想简单而引人注目：通过多个小步骤从噪声向量生成图像，将噪声逐步细化为潜在图像表示。
然而，这种方法不可避免地增加了总体推理时间，并在部署在客户机上时导致用户体验不佳。可以注意到，强大的GPU通常可以提供帮助，这是正确的，但是这样做的成本会大幅增加。以参考为例，在2023年上半年，具有8个vCPU和64 GB RAM的强大CPU r6i.2xlarge实例的价格为每小时0.504美元，而具有16 GB VRAM的类似g4dn.2xlarge实例与NVIDIA T4的价格为每小时0.75美元，即增加了1.5倍..
这使得图像生成服务对其所有者和用户来说相当昂贵。在运行在用户端的客户应用程序中，问题甚至更加严重。甚至可能没有GPU！这使得部署稳定扩散流水线成为一个具有挑战性的问题。
在过去的五年中，OpenVINO Toolkit封装了许多高性能推理功能。最初设计用于计算机视觉模型，它仍然在该领域占据主导地位，为许多当代模型提供了最佳推理性能，包括稳定扩散。然而，将稳定扩散模型优化为受资源限制的应用程序需要超越运行时优化。这正是OpenVINO神经网络压缩框架（NNCF）的模型优化能力发挥作用的地方。
在本博客文章中，我们将概述优化稳定扩散模型的问题，并提出一种工作流程，可以大大减少这些模型在运行在资源受限的硬件（如CPU）上时的延迟。特别地，与PyTorch相比，我们实现了5.1倍的推理加速和4倍的模型占用空间减少。
稳定扩散优化
在稳定扩散流水线中，UNet模型的计算成本最高。因此，优化一个模型在推理速度方面带来了显著的好处。
然而，事实证明，传统的模型优化方法，例如训练后的8位量化，对于这个模型不起作用。这主要有两个原因。首先，像语义分割，超分辨率等像素级预测模型是模型优化方面最复杂的之一，因为任务的复杂性导致调整模型参数和结构会以多种方式破坏结果。第二个原因是该模型具有较低的冗余级别，因为它在数亿个样本上进行训练时容纳了大量信息。这就是为什么研究人员必须采用更复杂的量化方法来在优化后保持准确性的原因。例如，Qualcomm使用了逐层知识蒸馏方法（AdaRound）来量化稳定扩散模型。这意味着无论如何都需要在量化后进行模型调整。如果是这样，为什么不使用量化感知训练（QAT），它可以同时调整模型和量化参数，以与源模型的训练方式相同？因此，我们在我们的工作中使用了NNCF，OpenVINO和Diffusers的这种方法，并将其与Token Merging相结合。
优化工作流程
通常，我们在模型训练完成后开始对模型进行优化。在这里，我们从在Pokemons数据集上进行了微调的模型开始，该数据集包含Pokemons的图像和它们的文本描述。
我们使用了Diffusers中的稳定扩散文本到图像微调示例，并将NNCF中的QAT集成到以下训练脚本中。我们还更改了损失函数，以将知识蒸馏从作为教师的源模型中引入到此过程中进行训练的实际模型中。这种方法不同于经典的知识蒸馏方法，其中经过训练的教师模型被蒸馏为较小的学生模型。在我们的情况下，知识蒸馏被用作辅助方法，有助于提高优化模型的最终准确性。我们还使用指数移动平均（EMA）方法来处理模型参数，但不包括量化器，这使得训练过程更加稳定。我们仅对模型进行了4096次迭代的调整。
通过一些技巧，例如梯度检查点和将EMA模型保存在RAM而不是VRAM中，我们可以使用一个具有24 GB VRAM的GPU运行优化过程。整个优化过程仅需要一天时间使用一个GPU完成！
超越量化感知训练
仅仅进行量化即可通过减少模型占用空间、加载时间、内存消耗和推理延迟来带来显著的改进。但是量化的好处在于可以与其他优化方法一起应用，从而实现累积加速。
最近，Facebook Research引入了一种用于Vision Transformer模型的Token Merging方法。该方法的核心是使用其中一种可用的策略（平均值、最大值等）将冗余的令牌与重要的令牌合并。这是在自注意力块之前完成的，自注意力块是Transformer模型中计算最复杂的部分。因此，减少令牌维度可以减少自注意力块的总体计算时间。该方法还已经用于Stable Diffusion模型，并且在优化在GPU上运行的用于高分辨率图像合成的Stable Diffusion流水线时，已经显示出有希望的结果。
我们修改了Token Merging方法以符合OpenVINO，并在应用于Attention UNet模型时与8位量化一起堆叠。这还涉及到所有提到的技术，包括知识蒸馏等。至于量化，需要进行微调才能恢复准确性。我们还从在Pokemons数据集上训练的模型开始优化和微调。下图显示了整体优化工作流程。
当在计算资源有限的设备上进行推理时，得到的模型非常有益，例如客户端或边缘CPU。正如前面提到的，将Token Merging与量化堆叠在一起会进一步减少推理延迟。
使用不同优化模型进行图像生成演示的结果。输入提示为“卡通鸟”，种子为42。模型使用Hugging Face Spaces中的OpenVINO 2022.3，在使用了第三代Intel® Xeon®可扩展处理器与Intel® Deep Learning Boost技术的“CPU升级”实例上运行。
结果
我们使用公开的优化工作流程得到了两种类型的优化模型，即8位量化和带有Token Merging的量化模型，并将它们与PyTorch基准进行了比较。我们还将基准转换为纯净的OpenVINO浮点（FP32）模型，以进行全面比较。
上图显示了图像生成的结果和一些模型特性。正如您所见，仅仅转换为OpenVINO就可以显著减少推理延迟（ 1.9倍 ）。应用8位量化进一步提高推理速度，相对于PyTorch可以实现3.9倍 的加速。量化的另一个好处是显著减少模型占用空间，仅为PyTorch检查点的 0.25倍 ，这也提高了模型加载时间。在量化的基础上应用Token Merging（ToME）（合并比例为 0.4 ）可以在保持模型占用空间不变的同时实现5.1倍 的性能加速。我们没有对优化模型的视觉质量进行详细分析，但是正如您所见，结果非常可靠。
下面我们展示了优化为在Intel CPU上运行的最终流水线的推理方法：
from optimum.intel import OVStableDiffusionPipeline
# 加载并编译性能流水线。
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()
# 生成图像。
prompt = "一只带红眼睛的绿色宝可梦的画"
output = pipe(prompt, num_inference_steps=50, output_type="pil").images[0]
output.save("image.png")您可以在Hugging Face Optimum Intel库中找到训练和量化代码。演示优化模型与原始模型之间差异的笔记本可在此处找到。您还可以在Hugging Face Hub的OpenVINO组织下找到许多模型。此外，我们还创建了一个在Hugging Face Spaces上运行的演示，该演示在r6id.2xlarge实例上使用第三代Intel Xeon可扩展处理器运行。
通用的稳定扩散模型怎么样？
正如我们在Pokemon图像生成任务中展示的那样，使用相对较少的训练资源可以实现Stable Diffusion流水线的高度优化。与此同时，众所周知，训练通用的稳定扩散模型是一项昂贵的任务。然而，如果有足够的预算和硬件资源，可以使用所描述的方法优化通用模型，并对其进行调整以生成高质量的图像。我们唯一的注意事项是与令牌合并方法相关的，它会大大降低模型容量。在这里的经验法则是，训练数据集越复杂，优化过程中使用的合并比例就应该越小。
如果您喜欢阅读这篇文章，您可能也对查看这篇文章感兴趣，该文章讨论了优化第四代英特尔Xeon CPU上稳定扩散性能的其他补充方法。
