使用🧨 Diffusers和Dreambooth訓練Stable Diffusion

釋出於 2022 年 11 月 7 日

在 GitHub 上更新

贊

訪客

Dreambooth 是一種使用專門的微調形式來教 Stable Diffusion 新概念的技術。有些人用它來將自己的照片置於奇幻場景中，而另一些人則用它來融入新的風格。🧨 Diffusers 提供了一個 Dreambooth 訓練指令碼。它的訓練時間不長，但很難選擇正確的超引數集，而且很容易過擬合。

我們進行了大量實驗，分析了 Dreambooth 中不同設定的影響。本文將介紹我們的發現，以及在使用 Dreambooth 微調 Stable Diffusion 時提高結果的一些技巧。

在開始之前，請注意，此方法絕不應用於惡意目的，以任何方式造成傷害，或未經他人知情而冒充他人。使用此方法訓練的模型仍受管理 Stable Diffusion 模型分發的 CreativeML Open RAIL-M 許可證約束。

注意：本文的早期版本已作為 W&B 報告釋出.

TL;DR: 推薦設定

Dreambooth 很容易過擬合。為了獲得高質量的影像，我們必須在訓練步數和學習率之間找到一個“最佳點”。我們建議使用較低的學習率，並逐步增加步數，直到結果令人滿意。
Dreambooth 需要更多訓練步數來處理人臉。在我們的實驗中，當批大小為 2 且學習率為 1e-6 時，800-1200 步效果很好。
在訓練人臉時，先驗保留對於避免過擬合很重要。對於其他主體，它似乎沒有太大區別。
如果您發現生成的影像有噪聲或質量下降，這很可能意味著過擬合。首先，嘗試上述步驟來避免它。如果生成的影像仍然有噪聲，請使用 DDIM 排程器或執行更多推理步數（在我們的實驗中，約 100 步效果很好）。
除了 UNet 之外，訓練文字編碼器對質量有很大影響。我們的最佳結果是透過文字編碼器微調、低學習率和適當的步數組合獲得的。然而，微調文字編碼器需要更多的記憶體，因此至少具有 24 GB RAM 的 GPU 是理想的。使用 8 位 Adam、`fp16` 訓練或梯度累積等技術，可以在 16 GB GPU（例如 Google Colab 或 Kaggle 提供的 GPU）上進行訓練。
使用或不使用 EMA 進行微調產生了相似的結果。
無需使用 `sks` 詞來訓練 Dreambooth。最初的實現之一使用它是因為它是詞彙表中一個罕見的詞元，但它實際上是一種步槍。我們的實驗以及例如 @nitrosocke 的實驗表明，選擇自然用於描述目標物件的術語是可以的。

學習率影響

Dreambooth 過擬合非常快。為了獲得好的結果，請根據您的資料集調整學習率和訓練步數。在我們的實驗中（詳見下文），我們對四個不同的資料集進行了高低學習率的微調。在所有情況下，我們都獲得了較低學習率下的更好結果。

實驗設定

我們所有的實驗均使用 train_dreambooth.py 指令碼和 AdamW 最佳化器在 2 個 40GB A100 GPU 上進行。我們使用相同的隨機種子，並保持所有超引數（除了學習率、訓練步數和是否使用先驗保留）在所有執行中相等。

對於前 3 個示例（各種物件），我們使用批大小 4（每個 GPU 2 個）對模型進行了 400 步的微調。我們使用了高學習率 5e-6 和低學習率 2e-6。沒有使用先驗保留。

最後一個實驗嘗試向模型中新增人類主體。在這種情況下，我們使用了批大小 2（每個 GPU 1 個）的先驗保留，並進行了 800 和 1200 步的訓練。我們使用了高學習率 5e-6 和低學習率 2e-6。

請注意，您可以使用 8 位 Adam、fp16 訓練或梯度累積來減少記憶體需求，並在具有 16 GB 記憶體的 GPU 上執行類似實驗。

貓玩具

高學習率（5e-6）

低學習率（2e-6）

豬頭

高學習率（5e-6）。請注意，顏色偽影是殘留的噪聲——執行更多推理步驟可能有助於解決這些細節。

低學習率（2e-6）

土豆頭先生

高學習率（5e-6）。請注意，顏色偽影是殘留的噪聲——執行更多推理步驟可能有助於解決這些細節。

低學習率（2e-6）

人臉

我們嘗試將《宋飛正傳》中的克萊默角色融入 Stable Diffusion。如前所述，我們以更小的批次大小訓練了更多步數。即便如此，結果也並不出色。為了簡潔起見，我們省略了這些樣本影像，並將讀者推遲到下一節，在那裡人臉訓練成為我們工作的重點。

初步結果總結

要使用 Dreambooth 訓練 Stable Diffusion 以獲得良好效果，調整學習率和訓練步數以適應您的資料集非常重要。

高學習率和過多的訓練步數會導致過擬合。模型將主要根據您的訓練資料生成影像，無論使用什麼提示。
低學習率和過少的步數會導致欠擬合：模型將無法生成我們試圖融入的概念。

人臉更難訓練。在我們的實驗中，對於物件，學習率為 2e-6 且訓練步數為 400 效果很好，但人臉需要 1e-6（或 2e-6）且約 1200 步。

如果模型過擬合，影像質量會嚴重下降，這種情況發生在：

學習率過高。
我們運行了過多的訓練步數。
在人臉的情況下，當不使用先驗保留時，如下一節所示。

訓練人臉時使用先驗保留

先驗保留是一種技術，它在微調過程中使用與我們正在嘗試訓練的類相同的附加影像。例如，如果我們嘗試將一個新人融入模型中，我們想要保留的*類*可能是*人*。先驗保留試圖透過結合新人的照片與其他人的照片來減少過擬合。好訊息是我們可以使用 Stable Diffusion 模型本身生成這些額外的類影像！如果您願意，訓練指令碼會自動處理，但您也可以提供一個包含您自己的先驗保留影像的資料夾。

先驗保留，1200 步，學習率=2e-6。