使用🧨 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。