使用 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 效能。
