用 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 位量化，并不适用于该模型。这主要有两个原因。首先，像素级预测模型（如语义分割、超分辨率等）由于任务复杂，是模型优化中最复杂的类型之一，因此调整模型参数和结构会以多种方式破坏结果。第二个原因是该模型的冗余度较低，因为它在数亿个样本上训练时容纳了大量信息。这就是为什么研究人员必须采用更复杂的量化方法，以在优化后保持准确率。例如，Qualcomm 使用逐层知识蒸馏方法（AdaRound）来量化 Stable Diffusion 模型。这意味着无论如何都需要在量化后对模型进行调优。既然如此，为什么不直接使用量化感知训练（QAT），以与训练源模型相同的方式同时调优模型和量化参数呢？因此，我们在工作中尝试了这种方法，使用了 NNCF、OpenVINO 和 Diffusers，并将其与Token Merging 结合。

优化工作流程

我们通常会在模型训练完成后开始优化。在这里，我们从一个在 Pokemons 数据集上微调过的模型开始，该数据集包含 Pokemons 的图像及其文本描述。

我们使用了 Diffusers 中 Stable Diffusion 的文本到图像微调示例，并将 NNCF 中的 QAT 集成到以下训练脚本中。我们还更改了损失函数，以纳入来自源模型的知识蒸馏，源模型在此过程中充当教师，而实际正在训练的模型充当学生。这种方法不同于经典的知识蒸馏方法，在经典方法中，训练好的教师模型被蒸馏到一个更小的学生模型中。在我们的案例中，知识蒸馏被用作一种辅助方法，有助于提高优化模型的最终准确率。我们还对模型参数使用指数移动平均（EMA）方法，但不包括量化器，这使我们能够使训练过程更加稳定。我们仅对模型进行了4096次迭代的调优。

借助一些技巧，例如梯度检查点和将 EMA 模型保留在 RAM 而不是 VRAM 中，我们可以使用一块具有 24 GB VRAM 的 GPU 运行优化过程。整个优化过程使用一块 GPU 耗时不到一天！

超越量化感知训练

仅量化就可以通过减少模型占用空间、加载时间、内存消耗和推理延迟带来显著提升。但量化的妙处在于，它可以与其他优化方法一起应用，从而实现累积加速。

最近，Facebook Research 为 Vision Transformer 模型引入了一种 Token Merging 方法。该方法的本质是使用可用策略之一（取平均、取最大值等）将冗余 token 与重要 token 合并。这在自注意力块之前完成，而自注意力块是 Transformer 模型中计算需求最高的部分。因此，减少 token 维度会减少自注意力块中的整体计算时间。该方法也已被适配用于 Stable Diffusion 模型，并在为 GPU 上运行的高分辨率图像合成优化 Stable Diffusion 流水线时展现出良好结果。

我们修改了 Token Merging 方法，使其兼容 OpenVINO，并在应用于 Attention UNet 模型时将其与 8 位量化叠加。这也涉及所有提到的技术，包括知识蒸馏等。至于量化，它需要进行微调以恢复准确率。我们还从在 Pokemons 数据集上训练的模型开始优化和微调。下图展示了整体优化工作流。

当在计算资源有限的设备上运行推理时，例如客户端或边缘 CPU，所得模型非常有益。如前所述，将 Token Merging 与量化叠加会进一步降低推理延迟。

使用不同优化模型进行图像生成演示的结果。输入提示词为“cartoon bird”，种子为42。这些模型使用 OpenVINO 2022.3，位于 HF Mirror Spaces 中，并使用“CPU upgrade”实例，该实例采用支持 Intel® Deep Learning Boost 技术的第3代英特尔® 至强® 可扩展处理器。

结果

我们使用公开的优化工作流程获得了两种优化模型：8 位量化模型和结合 Token Merging 的量化模型，并将它们与 PyTorch 基线进行比较。我们还把基线转换为原生 OpenVINO 浮点（FP32）模型，以进行全面的对比。

上图展示了图像生成的结果以及一些模型特性。可以看到，仅转换为 OpenVINO 就能显著降低推理延迟（1.9 倍）。应用 8 位量化进一步提升了推理速度，相比 PyTorch 实现了 3.9 倍的加速。量化的另一个好处是显著减小了模型体积，仅为 PyTorch checkpoint 的 0.25 倍，这也改善了模型加载时间。在量化的基础上应用 Token Merging（ToME）（合并比例为 0.4）带来了 5.1 倍的性能加速，同时保持体积不变。我们没有对优化模型的视觉质量进行详尽分析，但如你所见，结果相当可靠。

对于本博客中展示的结果，我们使用了默认的 50 个推理步数。推理步数越少，推理速度越快，但这会影响生成图像的质量。影响程度取决于模型和调度器。我们建议尝试不同的步数和调度器，找到最适合你的用例的方案。

下面我们展示如何使用最终优化为在英特尔 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")

你可以在 HF Mirror Optimum Intel 库中找到训练和量化代码。展示优化模型与原始模型差异的 notebook 可在此处获取。你还可以在 HF Mirror Hub 的 OpenVINO 组织下找到许多模型。此外，我们在 HF Mirror Spaces 上创建了一个演示，运行在第三代英特尔至强可扩展处理器上。

那通用 Stable Diffusion 模型呢？

正如我们在 Pokemon 图像生成任务中展示的那样，当使用相对较少的训练资源时，可以实现对 Stable Diffusion 流水线的高水平优化。同时，众所周知，训练通用 Stable Diffusion 模型是一项昂贵的任务。然而，如果有足够的预算和硬件资源，可以使用所述方法优化通用模型，并对其进行调优以生成高质量图像。我们唯一的注意事项与 token merging 方法有关，它会大幅降低模型容量。这里的经验法则是：训练数据集越复杂，优化过程中应使用的合并比例就越小。

如果你喜欢这篇文章，可能也会有兴趣查看这篇文章，其中讨论了在第四代英特尔至强 CPU 上优化 Stable Diffusion 性能的其他互补方法。