基于隐空间的扩散模型 (Latent Diffusion Model)，是解决文本到图片生成问题上的颠覆者。Stable Diffusion 是最驰名的一例，广泛应用在商业和工业。Stable Diffusion 的想法简略且无效: 从噪声向量开始，屡次去噪，以使之在隐空间里迫近图片的示意。
然而，这样的办法不可避免地减少了推理时长，使客户端的体验大打折扣。家喻户晓，一个好的 GPU 总能有帮忙，的确如此，但其损耗大大增加了。就推理而言，在 2023 年上半年 (H1’23)，一个好 CPU 实例 (r6i.2xlarge，8 vCPUs ，64 GB 内存) 价格是 0.504 $/h，同时，相似地，一个好 GPU 实例 ([g4dn.2xlarge](https://aws.amazon.com/ec2/instance-types/g4/)，NVIDIA T4，16 GB 内存) 价格是 0.75 $/h ，是前者的近 1.5 倍。
这就使图像生成的服务变得低廉，无论持有者还是用户。该问题在面向用户端部署就更突出了: 可能没有 GPU 能用！这让 Stable Diffusion 的部署变成了辣手的问题。
在过来五年中，OpenVINO 集成了许多高性能推理的个性。 其一开始为计算机视觉模型设计，现今仍在许多模型的推理性能上获得最佳体现，包含 Stable Diffusion。然而，对资源无限型的利用，优化 Stable Diffusion 远不止运行时的。这也是 OpenVINO NNCF(Neural Network Compression Framework) 发挥作用的中央。
在本博客中，咱们将理清优化 Stable Diffusion 模型的问题，并提出对资源无限的硬件 (比方 CPU) 减负的流程。尤其是和 PyTorch 相比，咱们速度进步了 5.1 倍，内存缩小了 4 倍。
Stable Diffusion 的优化
在 Stable Diffusion 的 管线 中，UNet 的运行是最计算低廉的。因而，对模型的推理速度，针对 UNet 的优化能带来足够的效益。
然而事实表明，传统的模型优化办法如 8-bit 的后训练量化，对此不见效。次要起因有两点: 其一，面向像素预测的模型，比方语义宰割、超分辨率等，是模型优化上最简单的，因为工作简单，参数和构造的扭转会导致无数种变数; 其二，模型的参数不是很冗余，因为其压缩了其数以千万计的 数据集 中的信息。这也是研究者不得不用更简单的量化办法来保障模型优化后的精度。举例而言，高通 (Qualcomm) 用分层常识蒸馏 (layer-wise Knowledge Distillation) 办法 (AdaRound) 来 量化 Stable Diffusion。这意味着，无论如何，模型量化后的微调是必要的。既然如此，为何不必 量化感知的训练 (Quantization-Aware Trainning, QAT)，其对原模型的微调和参数量化是同时进行的？因而，咱们在本工作中，用 token 合并 (Token Merging) 办法联合 NNCF, OpenVINO 和 Diffusers 实际了该想法。
优化流程
咱们通常从训练后的模型开始优化。在此，咱们从宝可梦数据集 (Pokemons dataset，蕴含图片和对应的文本形容) 上微调的 模型。
咱们对 Stable Diffusion 用 Diffusers 中的 图片 – 文本微调之例，联合 NNCF 中的 QAT (参见训练的 脚本)。咱们同时扭转了损失函数，以同时实现从源模型到部署模型的常识蒸馏。该办法与通常的常识蒸馏不同，后者是把源模型蒸馏到小些的模型。咱们的办法次要将常识整顿作为附加的办法，帮忙进步最初优化的模型的精度。咱们也用指数挪动均匀办法 (Exponential Moving Average, EMA) 让咱们训练过程更稳固。咱们仅对模型做 4096 次迭代。
基于一些技巧，比方梯度查看 (gradient checkpointing) 和 放弃 EMA 模型 在内存 (RAM) 而不是虚拟内存 (VRAM) 中。整个优化过程能用一张 GPU 在一天内实现。
量化感知的训练之外呢 ?
量化模型自身就能带来模型耗费、加载、内存、推理速度上的显著进步。但量化模型蛮大的劣势在能和其余模型优化办法一起，达到减速的增益成果。
最近，Facebook Research 针对视觉 Transformer 模型，提出了一个 Token Merging 办法。该办法的实质是用现有的办法 (取均匀、取最大值等) 把冗余的 token 和重要的 token 交融。这在 self-attention 块之前实现，后者是 Transformer 模型最耗费算力的局部。因而，减小 token 的跨度能缩小 self-attention 块耗费的工夫。该办法也已被 Stable Diffusion 模型 采纳，并在面向 GPU 的高分辨率优化上有可观的体现。
咱们改良了 Token Merging 办法，以便用 OpenVINO，并在注意力 UNet 模型上采纳 8-bit 量化。这蕴含了上述含常识蒸馏等的所有技术。对量化而言，其须要微调，以保障数值精度。咱们也从 宝可梦数据集 上训练的 模型 开始优化和微调。下图体现了总体的优化工作流程。
后果的模型在无限资源的硬件上是高度无效的，如客户机或边缘 CPU。如上文所述，把 Token Merging 办法和量化办法叠加能带来额定的推理增益。
<div class=”flex flex-row”>
<div class=”grid grid-cols-2 gap-4″>
<figure>
<img class=”max-w-full rounded-xl border-2 border-solid border-gray-600″ src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202307121519225.png” alt=”Image 1″ />
<figcaption class=”mt-2 text-center text-sm text-gray-500″>PyTorch FP32，推理时长：230.5 秒，内存耗费：3.44 GB</figcaption>
</figure>
<figure>
<img class=”max-w-full rounded-xl border-2 border-solid border-gray-600″ src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202307121519110.png” alt=”Image 2″ />
<figcaption class=”mt-2 text-center text-sm text-gray-500″>OpenVINO FP32，推理时长：120 秒 (1.9 倍)，内存耗费：3.44 GB</figcaption>
</figure>
<figure>
<img class=”max-w-full rounded-xl border-2 border-solid border-gray-600″ src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202307121519641.png” alt=”Image 3″ />
<figcaption class=”mt-2 text-center text-sm text-gray-500″>OpenVINO 8-bit，推理市场：59 秒 (3.9 倍)，内存耗费：0.86 GB(0.25 倍)</figcaption>
</figure>
<figure>
<img class=”max-w-full rounded-xl border-2 border-solid border-gray-600″ src=”https://man-archives.oss-cn-hangzhou.aliyuncs.com/goofan/202307121520495.png” alt=”Image 4″ />
<figcaption class=”mt-2 text-center text-sm text-gray-500″>ToMe + OpenVINO 8-bit, 推理速度：44.6 秒 (5.1 倍)，内存耗费：0.86 GB (0.25 倍)</figcaption>
</figure>
</div>
</div>
用不同模型优化办法的图片生成的后果 展现。输出提醒词为 “cartoon bird”，随机种子为 42。模型用 OpenVINO 2022.3，来自 Hugging Face Space，用“CPU 降级”的实例: 第三代 Intel® Xeon® Scalable Processors，和 Intel® 深度学习减速技术。
后果
咱们用优化模型不残缺的流程以失去两种模型: 基于 8-bit 量化的和基于 Token Merging 量化的，并和 PyTorch 作为基准比拟。咱们也把基准先转化成 vanilla OpenVINO (FP32) 的模型，以用以剖析性比拟。
下面的后果图展现了图像生成和局部模型的个性。如你所见，仅转化成 OpenVINO 就带来大的推理速度进步 ( 1.9 倍)。用基于 8-bit 的量化减速和 PyTorch 相比带来了 3.9 倍的推理速度。量化的另外一个重要进步在于内存耗费缩小，0.25 倍之于 PyTorch，同时也进步了加载速度。在量化之上利用 Token Merging (ToME) (交融比为 0.4) 带来了 5.1 倍 的提速，同时把模型内存耗费放弃在原程度上。咱们不提供输入后果上的品质扭转，但如你所见，后果还是有品质的。
上面咱们展现将最终优化后果部署在 Intel CPU 上代码。
from optimum.intel.openvino 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 库中你能够找到训练和量化 代码。比拟优化过的和原模型的 notebook 代码在 这里。你能够在 Hugging Face Hub 上找到 OpenVINO 下的 许多模型。另外，咱们在 Hugging Face Spaces 上建了一个 demo，以运行 [带第三代 Intel Xeon Scalable 的 r6id.2xlarge 实例]。
个别的 Stable Diffusion 模型呢？
正如咱们在宝可梦图像生成工作中展示的一样，仅用小量的训练资源，对 Stable Diffusion 管线实现高层次的优化是可能的。同时，家喻户晓，训练个别的 Stable Diffusion 模型是一个 低廉的工作。然而，有短缺的资金和硬件资源，用上述办法优化个别的模型生成高分辨率的模型是可能的。咱们惟一的正告是对于 Token Merging 办法，其会削弱模型容忍性。这里衡量标准是，训练数据越简单，优化模型时的交融比就该越小。
如果你乐于读本博客，那你可能对另外一篇 博客 感兴趣，它探讨了在第四代 Intel Xeon CPU 上其余互补的 Stable Diffusion 模型优化办法。
英文原文: https://hf.co/blog/train-optimize-sd-intel
作者: Alexander, Yury Gorbachev, Helena, Sayak Paul, Ella Charlaix
译者: Vermillion
审校/排版: zhongdongy (阿东)
发表回复