使用 NNCF 和 🤗 Optimum 优化 Intel CPU 上的 Stable Diffusion
潜在扩散模型在解决文本到图像生成问题方面是颠覆性的。 Stable Diffusion 是其中最著名的例子之一，在社区和行业中得到了广泛采用。Stable Diffusion 模型背后的思想简单而引人注目：您通过多个小步骤从噪声向量生成图像，将噪声细化为潜在图像表示。这种方法效果非常好，但如果您无法访问强大的 GPU，生成图像可能需要很长时间。
在过去的五年中，OpenVINO 工具包封装了许多用于高性能推理的功能。它最初是为计算机视觉模型设计的，至今仍在该领域占据主导地位，为许多当代模型（包括 Stable Diffusion）展示了同类最佳的推理性能。然而，为资源受限的应用程序优化 Stable Diffusion 模型需要超越运行时优化。这正是 OpenVINO 神经网络压缩框架 (NNCF) 的模型优化功能发挥作用的地方。
在这篇博文中，我们将概述优化 Stable Diffusion 模型的问题，并提出一种工作流程，该流程可大幅降低此类模型在 CPU 等资源受限硬件上运行时造成的延迟。具体来说，与 PyTorch 相比，我们实现了 5.1 倍的推理加速和 4 倍的模型占用空间减少。
Stable Diffusion 优化
在 Stable Diffusion 流水线中，UNet 模型是计算成本最高的部分。因此，仅优化一个模型就能在推理速度方面带来显著好处。
然而，事实证明，传统的模型优化方法，例如训练后 8 位量化，不适用于此模型。这主要有两个原因。首先，像素级预测模型，如语义分割、超分辨率等，由于任务的复杂性，在模型优化方面是最复杂的，因此调整模型参数和结构会以多种方式破坏结果。第二个原因是模型冗余度较低，因为它在 数亿个样本上进行训练时包含了大量信息。这就是为什么研究人员必须采用更复杂的量化方法来在优化后保持准确性。例如，高通公司使用层级知识蒸馏方法 (AdaRound) 来 量化 Stable Diffusion 模型。这意味着量化后仍然需要模型调优。既然如此，为什么不直接使用 量化感知训练 (QAT) 呢？它能够以与训练源模型相同的方式同时调优模型和量化参数。因此，我们在工作中使用 NNCF、OpenVINO 和 Diffusers 尝试了这种方法，并将其与 Token Merging 结合使用。
优化流程
我们通常在模型训练完成后开始优化。在这里，我们从一个在 宝可梦数据集（包含宝可梦图片及其文本描述）上微调的模型开始。
我们使用了 Diffusers 中的 文本到图像微调示例，并将 NNCF 中的 QAT 集成到以下训练脚本中。我们还更改了损失函数，以纳入来自源模型的知识蒸馏，该模型在此过程中充当教师，而实际训练的模型充当学生。这种方法与经典的知识蒸馏方法不同，后者将训练好的教师模型蒸馏到更小的学生模型中。在我们的案例中，知识蒸馏作为一种辅助方法，有助于提高优化模型的最终准确性。我们还对模型参数（不包括量化器）使用指数移动平均 (EMA) 方法，这使训练过程更加稳定。我们仅对模型进行了 4096 次迭代的调整。
通过一些技巧，例如梯度检查点和将 EMA 模型保留在 RAM 中而不是 VRAM 中，我们可以使用一块 24 GB VRAM 的 GPU 来运行优化过程。整个优化过程使用一块 GPU 不到一天即可完成！
超越量化感知训练
仅凭量化就能通过减少模型占用空间、加载时间、内存消耗和推理延迟来带来显著增强。但量化的好处在于它可以与其他优化方法一起应用，从而实现累积加速。
最近，Facebook 研究院推出了针对 Vision Transformer 模型的 Token Merging 方法。该方法的本质是利用可用的策略（平均、取最大值等）将冗余标记与重要标记合并。这在自注意力块之前完成，自注意力块是 Transformer 模型中计算要求最高的部分。因此，减少标记维度可以减少自注意力块中的总计算时间。该方法也已适用于 Stable Diffusion 模型，并在优化用于在 GPU 上运行高分辨率图像合成的 Stable Diffusion 管道时显示出有希望的结果。
我们修改了 Token Merging 方法，使其与 OpenVINO 兼容，并将其与 8 位量化结合，应用于 Attention UNet 模型。这还包括所有提到的技术，包括知识蒸馏等。至于量化，它需要微调才能应用以恢复准确性。我们还从在 宝可梦数据集上训练的模型开始优化和微调。下图显示了整体优化流程。
当在计算资源有限的设备（例如客户端或边缘 CPU）上运行推理时，所得到的模型非常有益。正如前面提到的，将 Token Merging 与量化相结合可以进一步减少推理延迟。
使用不同优化模型的图像生成演示结果。输入提示为“卡通鸟”，种子为 42。这些模型在 Hugging Face Spaces 中使用 OpenVINO 2022.3，并利用配备 Intel® 深度学习加速技术的第三代 Intel® Xeon® 可扩展处理器（“CPU 升级”实例）。
结果
我们使用了公开的优化流程，获得了两种类型的优化模型：8 位量化模型和使用 Token Merging 进行量化的模型，并将它们与 PyTorch 基线进行了比较。我们还将基线转换为香草 OpenVINO 浮点 (FP32) 模型以进行全面比较。
上图显示了图像生成的结果和一些模型特性。正如您所看到的，仅仅转换为 OpenVINO 就显着降低了推理延迟（1.9 倍）。应用 8 位量化进一步提升了推理速度，与 PyTorch 相比实现了 3.9 倍的加速。量化的另一个好处是模型占用空间显著减少，仅为 PyTorch 检查点的 0.25 倍，这也改善了模型加载时间。在量化之上应用 Token Merging (ToME)（合并率为 0.4）带来了 5.1 倍的性能提升，同时保持了相同的占用空间。我们没有对优化模型的视觉质量进行彻底分析，但是，如您所见，结果相当可靠。
对于本博客中显示的结果，我们使用了默认的 50 个推理步骤。推理步骤越少，推理速度越快，但这会影响生成图像的质量。这种影响的大小取决于模型和调度器。我们建议尝试不同数量的步骤和调度器，找出最适合您用例的方法。
下面我们展示了如何使用最终优化用于在 Intel CPU 上运行的管道进行推理
from optimum.intel import OVStableDiffusionPipeline
# Load and compile the pipeline for performance.
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()
# Generate an image.
prompt = "a drawing of a green pokemon with red eyes"
output = pipe(prompt, num_inference_steps=50, output_type="pil").images[0]
output.save("image.png")
您可以在 Hugging Face Optimum Intel 库中找到训练和量化代码。演示优化模型和原始模型之间差异的 Jupyter Notebook 可在此处获取。您还可以在 Hugging Face Hub 的 OpenVINO 组织下找到许多模型。此外，我们还在 Hugging Face Spaces 上创建了一个演示，该演示正在第三代 Intel Xeon 可扩展处理器上运行。
通用 Stable Diffusion 模型怎么样？
正如我们在宝可梦图像生成任务中所示，在训练资源相对较少的情况下，可以实现 Stable Diffusion 管道的高度优化。同时，众所周知，训练一个通用 Stable Diffusion 模型是一项昂贵的任务。然而，如果预算和硬件资源充足，可以使用所描述的方法优化通用模型并对其进行调整以生成高质量图像。我们唯一的注意事项与标记合并方法有关，该方法会大幅降低模型容量。经验法则是，训练数据集越复杂，优化过程中应使用的合并率越低。
如果您喜欢阅读这篇文章，您可能也会对查看这篇帖子感兴趣，其中讨论了其他补充方法来优化第四代 Intel Xeon CPU 上的 Stable Diffusion 性能。
