使用 LoRA 進行 Stable Diffusion 的高效引數微調

HuggingFace發表於2023-02-10

LoRA: Low-Rank Adaptation of Large Language Models 是微軟研究員引入的一項新技術，主要用於處理大模型微調的問題。目前超過數十億以上引數的具有強能力的大模型 (例如 GPT-3) 通常在為了適應其下游任務的微調中會呈現出巨大開銷。 LoRA 建議凍結預訓練模型的權重並在每個 Transformer 塊中注入可訓練層 (秩-分解矩陣)。因為不需要為大多數模型權重計算梯度，所以大大減少了需要訓練引數的數量並且降低了 GPU 的記憶體要求。研究人員發現，透過聚焦大模型的 Transformer 注意力塊，使用 LoRA 進行的微調質量與全模型微調相當，同時速度更快且需要更少的計算。

用於 Diffusers 的 LoRA ?

儘管 LoRA 最初是為大模型提出的，並在 transformer 塊上進行了演示，但該技術也可以應用於其他地方。在微調 Stable Diffusion 的情況下，LoRA 可以應用於將影像表示與描述它們的提示相關聯的交叉注意層。下圖的細節 (摘自 Stable Diffusion 論文) 並不重要，只需要注意黃色塊是負責建立圖文之間的關係表示就行。

潛在擴散架構

據我們所知，Simo Ryu (GitHub 使用者名稱 @cloneofsimo) 是第一個提出適用於 Stable Diffusion 的 LoRA 實現的人。如果想檢視相關示例和許多其他有趣的討論和見解。請一定要看看 他們的 GitHub 專案。

為了將 LoRA 的可訓練矩陣注入到與交叉注意力層一樣深的模型中，過去人們需要以富有想象力 (但脆弱) 的方式破解 diffusers 的原始碼。如果 Stable Diffusion 向我們展示了一件事，那就是社群總是會想出辦法來改變和調整模型以達到創造性目的，我們喜歡這樣！由於許多其他原因，提供操縱交叉注意力層的靈活性可能是有益的，例如更容易採用 xFormers 等最佳化技術。 Prompt-to-Prompt 等其他創意專案可以使用一些簡單的方法來訪問這些層，因此我們決定 為使用者提供一種通用的方法來做到這一點。自 12 月下旬以來，我們一直在測試，並在我們的 diffusers 中正式釋出。

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

  • 更快的訓練速度
  • 計算要求較低。我們可以在具有 11 GB VRAM 的 2080 Ti 中建立一個全微調模型！
  • 小了很多的訓練模型。由於原始模型已凍結，我們注入了新層進行訓練，因此我們可以將新層的權重儲存為大小約為 3 MB 的單個檔案。這比 UNet 模型的原始大小小一千倍

我們對最後一點特別興奮。為了讓使用者分享他們出色的微調或 dreamboothed 模型，他們必須分享最終模型的完整副本。其他想要試用它們的使用者必須在他們最喜歡的 UI 中下載經過微調的權重，這會增加大量儲存和下載成本。截至今天，大約有 1,000 個 Dreambooth 模型在 Dreambooth 概念庫中註冊，可能還有更多未在庫中註冊。

使用 LoRA，現在可以釋出 單個 3.29 MB 檔案 以允許其他人使用你的微調模型。

(感謝 GitHub 使用者 @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 dashboard，在 2080 Ti GPU (11 GB 記憶體)。我沒有嘗試最佳化超引數，所以請自行嘗試！Sayak 在 T4 (16 GB 記憶體) 上又跑了一次，這是 他的最終模型，這裡是 使用它的演示空間。

Sayak 的 LoRA 模型輸出示例

有關 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 選項，我們在上一節中看到的微調指令碼會自動填充有關基本模型的資訊。正如您在 pokemon-lora 的介紹文件 中所見，這被記錄為模型儲存庫的“自述檔案”檔案中的後設資料標籤。

在我們確定了用於使用 LoRA 進行微調的基礎模型之後，我們載入了一個正常的穩定擴散管道。我們將使用 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 權重 在常規模型權重之上，將 pipline 移動到 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 進行 Dreamboothing

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

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

要使用 LoRA 訓練 Dreambooth，您需要使用 此 diffusers 指令碼。請看一下 README、文件 和 我們的超引數探索博文 瞭解詳細資訊.

其他方法

對輕鬆微調的追求並不新鮮。除了 Dreambooth 之外，textual inversion 是另一種流行的方法，它試圖向訓練有素的穩定擴散模型教授新概念。使用 Textual Inversion 的主要原因之一是經過訓練的權重也很小且易於共享。然而，它們只適用於單個主題 (或一小部分主題)，而 LoRA 可用於通用微調，這意味著它可以適應新的領域或資料集。

Pivotal Tuning 是一種嘗試將 Textual Inversion 與 LoRA 相結合的方法。首先，您使用 textual inversion 技術向模型教授一個新概念，獲得一個新的標記嵌入來表示它。然後，您使用 LoRA 訓練該 token 嵌入以獲得兩全其美。

我們還沒有使用 LoRA 探索過 Pivotal Tuning。歡迎挑戰？ ?

英文原文: https://huggingface.co/blog/lora

原文作者: Pedro Cuenca, Sayak Paul

中文譯者: innovation64 (李洋)

審校: zhongdongy (忠東)

相關文章