基于隱空間的擴散模型 (Latent Diffusion Model)，是解決文本到圖片生成問題上的顛覆者。Stable Diffusion 是最著名的一例，廣泛應(yīng)用在商業(yè)和工業(yè)。Stable Diffusion 的想法簡單且有效: 從噪聲向量開始，多次去噪，以使之在隱空間里逼近圖片的表示。

但是，這樣的方法不可避免地增加了推理時長，使客戶端的體驗大打折扣。眾所周知，一個好的 GPU 總能有幫助，確實如此，但其損耗大大增加了。就推理而言，在 2023 年上半年 (H1’23)，一個好 CPU 實例 (r6i.2xlarge，8 vCPUs ，64 GB 內(nèi)存) 價格是 0.504 $/h，同時，類似地，一個好 GPU 實例 (g4dn.2xlarge，NVIDIA T4，16 GB 內(nèi)存) 價格是 0.75 $/h ，是前者的近 1.5 倍。

這就使圖像生成的服務(wù)變得昂貴，無論持有者還是用戶。該問題在面向用戶端部署就更突出了: 可能沒有 GPU 能用！這讓 Stable Diffusion 的部署變成了棘手的問題。

在過去五年中，OpenVINO 集成了許多高性能推理的特性。其一開始為計算機視覺模型設(shè)計，現(xiàn)今仍在許多模型的推理性能上取得最佳表現(xiàn)，包括 Stable Diffusion。然而，對資源有限型的應(yīng)用，優(yōu)化 Stable Diffusion 遠不止運行時的。這也是 OpenVINO NNCF(Neural Network Compression Framework) 發(fā)揮作用的地方。

在本博客中，我們將理清優(yōu)化 Stable Diffusion 模型的問題，并提出對資源有限的硬件 (比如 CPU) 減負的流程。尤其是和 PyTorch 相比，我們速度提高了 5.1 倍，內(nèi)存減少了 4 倍。

Stable Diffusion 的優(yōu)化

在 Stable Diffusion 的 管線 中，UNet 的運行是最計算昂貴的。因此，對模型的推理速度，針對 UNet 的優(yōu)化能帶來足夠的效益。

然而事實表明，傳統(tǒng)的模型優(yōu)化方法如 8-bit 的后訓(xùn)練量化，對此不奏效。主要原因有兩點: 其一，面向像素預(yù)測的模型，比如語義分割、超分辨率等，是模型優(yōu)化上最復(fù)雜的，因為任務(wù)復(fù)雜，參數(shù)和結(jié)構(gòu)的改變會導(dǎo)致無數(shù)種變數(shù); 其二，模型的參數(shù)不是很冗余，因為其壓縮了其數(shù)以千萬計的 數(shù)據(jù)集 中的信息。這也是研究者不得不用更復(fù)雜的量化方法來保證模型優(yōu)化后的精度。舉例而言，高通 (Qualcomm) 用分層知識蒸餾 (layer-wise Knowledge Distillation) 方法 (AdaRound) 來 量化 Stable Diffusion。這意味著，無論如何，模型量化后的微調(diào)是必要的。既然如此，為何不用 量化感知的訓(xùn)練 (Quantization-Aware Trainning, QAT)，其對原模型的微調(diào)和參數(shù)量化是同時進行的？因此，我們在本工作中，用 token 合并 (Token Merging) 方法結(jié)合 NNCF, OpenVINO 和 Diffusers 實踐了該想法。

優(yōu)化流程

我們通常從訓(xùn)練后的模型開始優(yōu)化。在此，我們從寶可夢數(shù)據(jù)集 (Pokemons dataset，包含圖片和對應(yīng)的文本描述) 上微調(diào)的 模型。

我們對 Stable Diffusion 用 Diffusers 中的 圖片 - 文本微調(diào)之例，結(jié)合 NNCF 中的 QAT (參見訓(xùn)練的 腳本)。我們同時改變了損失函數(shù)，以同時實現(xiàn)從源模型到部署模型的知識蒸餾。該方法與通常的知識蒸餾不同，后者是把源模型蒸餾到小些的模型。我們的方法主要將知識整理作為附加的方法，幫助提高最后優(yōu)化的模型的精度。我們也用指數(shù)移動平均方法 (Exponential Moving Average, EMA) 讓我們訓(xùn)練過程更穩(wěn)定。我們僅對模型做 4096 次迭代。

基于一些技巧，比如梯度檢查 (gradient checkpointing) 和 保持 EMA 模型 在內(nèi)存 (RAM) 而不是虛擬內(nèi)存 (VRAM) 中。整個優(yōu)化過程能用一張 GPU 在一天內(nèi)完成。

量化感知的訓(xùn)練之外呢 ?

量化模型本身就能帶來模型消耗、加載、內(nèi)存、推理速度上的顯著提高。但量化模型蠻大的優(yōu)勢在能和其他模型優(yōu)化方法一起，達到加速的增益效果。

最近，F(xiàn)acebook Research 針對視覺 Transformer 模型，提出了一個 Token Merging 方法。該方法的本質(zhì)是用現(xiàn)有的方法 (取平均、取最大值等) 把冗余的 token 和重要的 token 融合。這在 self-attention 塊之前完成，后者是 Transformer 模型最消耗算力的部分。因此，減小 token 的跨度能減少 self-attention 塊消耗的時間。該方法也已被 Stable Diffusion 模型 采用，并在面向 GPU 的高分辨率優(yōu)化上有可觀的表現(xiàn)。

我們改進了 Token Merging 方法，以便用 OpenVINO，并在注意力 UNet 模型上采用 8-bit 量化。這包含了上述含知識蒸餾等的所有技術(shù)。對量化而言，其需要微調(diào)，以保證數(shù)值精度。我們也從 寶可夢數(shù)據(jù)集 上訓(xùn)練的 模型 開始優(yōu)化和微調(diào)。下圖體現(xiàn)了總體的優(yōu)化工作流程。

結(jié)果的模型在有限資源的硬件上是高度有效的，如客戶機或邊緣 CPU。如上文所述，把 Token Merging 方法和量化方法疊加能帶來額外的推理增益。

用不同模型優(yōu)化方法的圖片生成的結(jié)果 展示。輸入提示詞為 “cartoon bird”，隨機種子為 42。模型用 OpenVINO 2022.3，來自 Hugging Face Space，用“CPU 升級”的實例: 第三代 Intel? Xeon? Scalable Processors，和 Intel? 深度學(xué)習(xí)加速技術(shù)。

結(jié)果

我們用優(yōu)化模型不完整的流程以得到兩種模型: 基于 8-bit 量化的和基于 Token Merging 量化的，并和 PyTorch 作為基準(zhǔn)比較。我們也把基準(zhǔn)先轉(zhuǎn)化成 vanilla OpenVINO (FP32) 的模型，以用以分析性比較。

上面的結(jié)果圖展示了圖像生成和部分模型的特性。如你所見，僅轉(zhuǎn)化成 OpenVINO 就帶來大的推理速度提高 (?1.9?倍)。用基于 8-bit 的量化加速和 PyTorch 相比帶來了?3.9?倍的推理速度。量化的另外一個重要提高在于內(nèi)存消耗減少，0.25 倍之于 PyTorch，同時也提高了加載速度。在量化之上應(yīng)用 Token Merging (ToME) (融合比為 0.4) 帶來了?5.1 倍?的提速，同時把模型內(nèi)存消耗保持在原水平上。我們不提供輸出結(jié)果上的質(zhì)量改變，但如你所見，結(jié)果還是有質(zhì)量的。

下面我們展示將最終優(yōu)化結(jié)果部署在 Intel CPU 上代碼。

在 Hugging Face Optimum Intel 庫中你可以找到訓(xùn)練和量化 代碼。比較優(yōu)化過的和原模型的 notebook 代碼在 這里。你可以在 Hugging Face Hub 上找到 OpenVINO 下的 許多模型。另外，我們在 Hugging Face Spaces 上建了一個 demo，以運行帶第三代 Intel Xeon Scalable 的 r6id.2xlarge 實例。

一般的 Stable Diffusion 模型呢？

正如我們在寶可夢圖像生成任務(wù)中展現(xiàn)的一樣，僅用小量的訓(xùn)練資源，對 Stable Diffusion 管線實現(xiàn)高層次的優(yōu)化是可能的。同時，眾所周知，訓(xùn)練一般的 Stable Diffusion 模型是一個 昂貴的任務(wù)。但是，有充足的資金和硬件資源，用上述方法優(yōu)化一般的模型生成高分辨率的模型是可能的。我們唯一的警告是關(guān)于 Token Merging 方法，其會減弱模型容忍性。這里衡量標(biāo)準(zhǔn)是，訓(xùn)練數(shù)據(jù)越復(fù)雜，優(yōu)化模型時的融合比就該越小。

如果你樂于讀本博客，那你可能對另外一篇?博客?感興趣，它討論了在第四代 Intel Xeon CPU 上其他互補的 Stable Diffusion 模型優(yōu)化方法。

?? 寶子們可以戳?閱讀原文?查看文中所有的外部鏈接喲！

英文原文:?https://hf.co/blog/train-optimize-sd-intel

作者: Alexander, Yury Gorbachev, Helena, Sayak Paul, Ella Charlaix

譯者: Vermillion

審校/排版: zhongdongy (阿東)