11 個月前 Stable Diffusion 誕生，它能夠在消費級 GPU 上運行的消息讓不少研究者備受鼓舞。不僅如此，蘋果官方很快下場，將 Stable Diffusion「塞進」iPhone、iPad 和 Mac 中運行。這大大降低了 Stable Diffusion 對硬件設(shè)備的要求，讓其逐漸成為人人都能使用的「黑科技」。

現(xiàn)在，它甚至已經(jīng)可以在 Raspberry Pi Zero 2 上運行了。

圖片

Raspberry Pi Zero 2 「Just as small. Five times as fast.」

這是怎樣一個概念？運行 Stable Diffusion 并不是一件容易的事，它包含一個 10 億參數(shù)的大型 Transformer 模型，建議使用的最低 RAM/VRAM 通常為 8GB。而 RPI Zero 2 只是內(nèi)存為 512MB 的微型計算機。

這意味著在 RPI Zero 2 上運行 Stable Diffusion 是一個巨大的挑戰(zhàn)。而且，在運行過程中，作者沒有增加存儲空間，也沒有將中間結(jié)果卸載到磁盤上。

一般而言，主要的機器學(xué)習框架和庫都專注于最小化推理延遲和 / 或最大化吞吐量，但以上這些都以內(nèi)存使用為代價。因此，作者決定寫一個超小的、可解開的推理庫，致力于將內(nèi)存消耗最小化。

OnnxStream 做到了。

項目地址 https://github.com/vitoplantamura/OnnxStream

OnnxStream 基于將推理引擎與負責提供模型權(quán)重的組件解耦的思路，后者是派生自 WeightsProvider 的一個類。一個 WeightsProvider 的專門化可以實現(xiàn)任何類型的模型參數(shù)加載、緩存和預(yù)取。例如，一個自定義的 WeightsProvider 可以決定直接從 HTTP 服務(wù)器下載數(shù)據(jù)，而不加載或?qū)懭肴魏蝺?nèi)容到磁盤（這也是 OnnxStream 命名中有 Stream 的原因）。有兩個默認的 WeightsProviders 可用：DiskNoCache 和 DiskPrefetch。

與微軟的推理框架 OnnxStream 相比，OnnxStream 只需要消耗 1/55 的內(nèi)存就可以達到同樣的效果，但（在 CPU 上的）速度只比前者慢 0.5-2 倍。

接下來你將看到 Stable Diffusion 在 RPI Zero 2 上運行的效果，以及背后的方法。需要注意的是，雖然運行速度較慢，但是它是大模型在更小、更有限的設(shè)備上運行的嶄新嘗試。

圖片

 網(wǎng)友們認為這個項目很酷

將 Stable Diffusion 在 Raspberry Pi Zero 2 上運行

VAE 解碼器是 Stable Diffusion 中唯一無法以單精度或半精度放入 RPI Zero 2 RAM 的模型。這是因為模型中存在殘差連接、非常大的張量和卷積。唯一的解決辦法就是靜態(tài)量化（8 bit）。

以下這些圖像是由作者 repo 中包含的 Stable Diffusion 示例實現(xiàn)在不同精度的 VAE 解碼器下使用 OnnxStream 生成的。

第一張圖像是在作者的 PC 上生成的，使用了由 RPI Zero 2 生成的相同的 latent。

精度為 W16A16 的 VAE 解碼器的生成效果

精度為 W8A32 的 VAE 解碼器的生成效果

第三張圖由 RPI Zero 2 在大約 3 小時內(nèi)生成。

圖注：精度為 W8A8 的 VAE 解碼器的生成效果

OnnxStream 的特點

• 推理引擎與 WeightsProvider 解耦
• WeightsProvider 可以是 DiskNoCache、DiskPrefetch 或自定義
• 注意力切片
• 動態(tài)量化（8 bit 無符號、非對稱、百分位數(shù)）
• 靜態(tài)量化（W8A8 無符號、非對稱、百分位數(shù)）
• 輕松校準量化模型
• 支持 FP16（使用或不使用 FP16 運算）
• 實現(xiàn)了 24 個 ONNX 算子（最常用的算子）
• 運算按順序執(zhí)行，但所有算子都是多線程的
• 單一實現(xiàn)文件 + header 文件
• XNNPACK 調(diào)用被封裝在 XnnPack 類中 (用于將來的替換)

并且需要注意的是，OnnxStream 依賴 XNNPACK 來加速某些原語：MatMul、Convolution、element-wise Add/Sub/Mul/Div、Sigmoid 和 Softmax。

性能對比

Stable Diffusion 由三個模型組成：文本編碼器（672 次運算和 1.23 億個參數(shù)）、UNET 模型（2050 次運算和 8.54 億個參數(shù)）和 VAE 解碼器（276 次運算和 4900 萬個參數(shù)。

假設(shè)批大小等于 1，生成完整圖像則需要 10 步，這需要運行 2 次文本編碼器、運行 20 次（即 2*10）UNET 模型和運行 1 次 VAE 解碼器，才能獲得良好效果（使用 Euler Ancestral 調(diào)度器）。

該表顯示了 Stable Diffusion 的三個模型不同的推理時間，以及內(nèi)存消耗（即 Windows 中的 Peak Working Set Size 或 Linux 中的 Maximum Resident Set Size）。

可以發(fā)現(xiàn)，在 UNET 模型中（以 FP16 精度運行時，OnnxStream 中啟用了 FP16 算術(shù)），OnnxStream 的內(nèi)存消耗量僅為 OnnxRuntime 的 1/55，但速度只慢 0.5-2 倍。

這次測試需要注明的幾點是：

• OnnxRuntime 的第一次運行是預(yù)熱推理，因為它的 InferenceSession 是在第一次運行前創(chuàng)建的，并在隨后的所有運行中重復(fù)使用。而 OnnxStream 沒有預(yù)熱推理，因為它的設(shè)計是純粹「eager」的（不過，后續(xù)運行可以受益于操作系統(tǒng)對權(quán)重文件的緩存）。
• 目前 OnnxStream 不支持 batch size ! = 1 的輸入，這與 OnnxRuntime 不同，后者在運行 UNET 模型時使用 batch size = 2 可以大大加快整個擴散過程。
• 在測試中，改變 OnnxRuntime 的 SessionOptions（如 EnableCpuMemArena 和 ExecutionMode）對結(jié)果沒有產(chǎn)生明顯影響。
• 在內(nèi)存消耗和推理時間方面，OnnxRuntime 的性能與 NCNN（另一個框架）非常相似。
• 測試的運行條件：Windows Server 2019、16GB 內(nèi)存、8750H CPU (AVX2)、970 EVO Plus SSD, VMWare 上的 8 個虛擬內(nèi)核。

注意力切片與量化

在運行 UNET 模型時，采用「注意力切片」技術(shù)，并對 VAE 解碼器使用 W8A8 量化，這對于將模型內(nèi)存消耗降低到適合在 RPI Zero 2 上運行的水平至關(guān)重要。

雖然互聯(lián)網(wǎng)上有很多關(guān)于量化神經(jīng)網(wǎng)絡(luò)的信息，但關(guān)于「注意力切片」的卻很少。

這里的想法很簡單：目標是在計算 UNET 模型中各種多頭注意力的縮放點積注意力時，避免生成完整的 Q @ K^T 矩陣。在 UNET 模型中，注意力頭數(shù)為 8 時，Q 的形狀為 (8,4096,40)，同時 K^T 為 (8,40,4096)。因此，第一個 MatMul 的最終形狀為 (8,4096,4096)，這是一個 512MB 的張量（FP32 精度）。

圖片

解決方案是垂直分割 Q，然后在每個 Q 塊上正常進行注意力操作。Q_sliced 形狀為 (1,x,40)，其中 x 為 4096（在本例中），除以 onnxstream::Model::m_attention_fused_ops_parts（默認值為 2，但可以自定義。

這個簡單的技巧可以將 UNET 模型以 FP32 精度運行時的整體內(nèi)存消耗從 1.1GB 降低到 300MB。一個更高效的替代方案是使用 FlashAttention，但是 FlashAttention 需要為每個支持的架構(gòu)（AVX, NEON）等編寫自定義內(nèi)核，在作者給出的例子中繞過 XnnPack。

更多信息參見該項目的 GitHub 界面。