使用 LoRA 高效微調 Stable Diffusion

LoRA: Low-Rank Adaptation of Large Language Models (LoRA：大型語言模型的低秩適應) 是微軟研究人員引入的一項新技術，用於解決微調大型語言模型的問題。像 GPT-3 這樣擁有數十億引數的強大模型，其微調成本極其高昂，難以適應特定任務或領域。LoRA 提議凍結預訓練模型的權重，並在每個 Transformer 塊中注入可訓練的層（秩分解矩陣）。這極大地減少了可訓練引數的數量和 GPU 記憶體需求，因為大部分模型權重無需計算梯度。研究人員發現，透過專注於大型語言模型的 Transformer 注意力模組，使用 LoRA 的微調質量與全模型微調相當，同時速度更快，計算需求更少。

LoRA for Diffusers 🧨

儘管 LoRA 最初是為大型語言模型提出的，並在 Transformer 模組上得到了驗證，但該技術也可以應用於其他地方。在 Stable Diffusion 微調的情況下，LoRA 可以應用於將影像表示與描述它們的提示詞相關聯的交叉注意力層。下圖（取自 Stable Diffusion 論文）的細節不重要，只需注意黃色模組是負責建立影像和文字表示之間關係的。

據我們所知，Simo Ryu (@cloneofsimo) 是第一個提出適用於 Stable Diffusion 的 LoRA 實現的人。請務必檢視他們的 GitHub 專案，以瞭解示例和許多有趣的討論和見解。

為了將 LoRA 可訓練矩陣注入到模型深處的交叉注意力層，人們過去需要以富有想象力（但不穩定）的方式修改 diffusers 的原始碼。如果說 Stable Diffusion 教會了我們一件事，那就是社群總能想出辦法來改造和調整模型以實現創意目的，我們對此非常喜愛！提供操縱交叉注意力層的靈活性可能在許多其他方面都有益，例如更容易採用像 xFormers 這樣的最佳化技術。其他創意專案，如 Prompt-to-Prompt，也可以透過某種簡單的方式訪問這些層，因此我們決定為使用者提供一種通用的方式來實現它。自 12 月下旬以來，我們一直在測試那個 *pull request*，它已於昨天隨我們的 diffusers 版本正式釋出。

我們一直與 @cloneofsimo 合作，在 diffusers 中為 Dreambooth 和全微調方法提供 LoRA 訓練支援！這些技術提供了以下好處：

• 如前所述，訓練速度要快得多。
• 計算要求更低。我們可以在一塊 11 GB 視訊記憶體的 2080 Ti 上建立一個完整的微調模型！
• 訓練出的權重非常非常小。由於原始模型被凍結，我們注入新的層進行訓練，我們可以將新層的權重儲存為一個大小約為 3 MB 的檔案。這比 UNet 模型的原始大小小了約*一千倍*！

我們對最後一點尤其感到興奮。為了讓使用者分享他們出色的微調或 *dreambooth* 模型，他們必須分享最終模型的完整副本。其他想嘗試這些模型的使用者必須在他們喜歡的 UI 中下載微調後的權重，這導致了巨大的儲存和下載成本。截至今天，在 Dreambooth Concepts Library 中註冊了大約 1,000 個 Dreambooth 模型，可能還有更多未在庫中註冊的模型。

有了 LoRA，現在可以釋出一個 3.29 MB 的檔案，讓其他人使用你微調後的模型。

（感謝 @mishig25，我第一個聽到他在日常對話中把 **dreamboothing** 當作動詞使用的人）。

LoRA 微調

Stable Diffusion 的全模型微調過去既慢又困難，這也是為什麼像 Dreambooth 或 Textual Inversion 這樣更輕量級的方法變得如此流行的部分原因。有了 LoRA，在自定義資料集上微調模型變得容易得多。

Diffusers 現在提供了一個 LoRA 微調指令碼，它可以在低至 11 GB 的 GPU RAM 中執行，而無需藉助 8-bit 最佳化器等技巧。以下是如何使用它來微調一個使用 Lambda Labs Pokémon 資料集 的模型：

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/pokemon"
export HUB_MODEL_ID="pokemon-lora"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --dataloader_num_workers=8 \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-04 \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" --lr_warmup_steps=0 \
  --output_dir=${OUTPUT_DIR} \
  --push_to_hub \
  --hub_model_id=${HUB_MODEL_ID} \
  --report_to=wandb \
  --checkpointing_steps=500 \
  --validation_prompt="Totoro" \
  --seed=1337

值得注意的是，學習率是 1e-4，遠大於常規微調的學習率（通常在 ~1e-6 的數量級）。這是之前執行的 W&B 儀表板，在 2080 Ti GPU (11 GB RAM) 上花費了大約 5 小時。我沒有嘗試最佳化超引數，所以請隨意自己嘗試！Sayak 在 T4 (16 GB RAM) 上進行了另一次執行，這是他的最終模型，這裡是一個使用它的演示 Space。

有關 diffusers 中 LoRA 支援的更多詳細資訊，請參閱我們的文件——它將始終與實現保持同步更新。

推理

正如我們所討論的，LoRA 的一個主要優勢是，透過訓練比原始模型大小少數個數量級的權重，就能獲得出色的結果。我們設計了一個推理流程，允許在未經修改的 Stable Diffusion 模型權重之上載入額外的權重。讓我們看看它是如何工作的。

首先，我們將使用 Hub API 自動確定用於微調 LoRA 模型的基礎模型是什麼。從 Sayak 的模型開始，我們可以使用以下程式碼：

from huggingface_hub import model_info

# LoRA weights ~3 MB
model_path = "sayakpaul/sd-model-finetuned-lora-t4"

info = model_info(model_path)
model_base = info.cardData["base_model"]
print(model_base)   # CompVis/stable-diffusion-v1-4

這段程式碼將打印出他用於微調的模型，即 CompVis/stable-diffusion-v1-4。在我的情況下，我從 Stable Diffusion 的 1.5 版本開始訓練我的模型，所以如果你用我的 LoRA 模型執行相同的程式碼，你會看到輸出是 runwayml/stable-diffusion-v1-5。

如果你使用 --push_to_hub 選項，我們上一節看到的微調指令碼會自動填充有關基礎模型的資訊。這會作為元資料標籤記錄在模型倉庫的 README 檔案中，你可以在這裡看到。

在我們確定了用於 LoRA 微調的基礎模型之後，我們載入一個正常的 Stable Diffusion 流水線。我們將使用 DPMSolverMultistepScheduler 對其進行自定義，以實現非常快的推理：

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

神奇之處就在這裡。我們在常規模型權重之上從 Hub 載入 LoRA 權重，將流水線移至 cuda 裝置並執行推理：

pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

image = pipe("Green pokemon with menacing face", num_inference_steps=25).images[0]
image.save("green_pokemon.png")

使用 LoRA 進行 Dreambooth

Dreambooth 允許你向 Stable Diffusion 模型“教授”新概念。LoRA 與 Dreambooth 相容，其過程與微調類似，但有幾個優點：

• 訓練更快。
• 我們只需要少量我們想要訓練的主體的影像（通常 5 或 10 張就足夠了）。
• 如果需要，我們可以調整文字編碼器，以提高對主體的保真度。

要使用 LoRA 訓練 Dreambooth，你需要使用這個 diffusers 指令碼。請檢視 README、文件和我們關於超引數探索的部落格文章以獲取詳細資訊。

要以一種快速、廉價且簡單的方式使用 LoRA 訓練你的 Dreambooth 模型，請檢視 hysts 的這個 Space。你需要複製它並分配一個 GPU 以便快速執行。這個過程將使你免於設定自己的訓練環境，並且你將能夠在幾分鐘內訓練你的模型！

其他方法

對簡單微調的追求並不新鮮。除了 Dreambooth，Textual Inversion 是另一種流行的方法，它試圖向訓練好的 Stable Diffusion 模型教授新概念。使用 Textual Inversion 的主要原因之一是訓練後的權重也很小且易於分享。然而，它們只適用於單個主體（或少數幾個），而 LoRA 可用於通用目的的微調，這意味著它可以適應新的領域或資料集。

Pivotal Tuning 是一種試圖將 Textual Inversion 與 LoRA 結合起來的方法。首先，你使用 Textual Inversion 技術教模型一個新概念，獲得一個新的詞嵌入來表示它。然後，你使用 LoRA 訓練該詞嵌入，以獲得兩者的優點。

我們還沒有探索過 Pivotal Tuning 與 LoRA 的結合。誰願意接受挑戰？🤗