使用 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 的结合。谁愿意接受挑战？🤗