使用 LoRA 进行高效稳定的扩散微调

LoRA: Low-Rank Adaptation of Large Language Models是微软研究人员为处理微调大型语言模型的问题而引入的一项新技术。具有数十亿参数的强大模型（例如 GPT-3）为了使其适应特定任务或领域而进行微调的成本高得令人望而却步。LoRA 提出冻结预训练模型权重并注入可训练层（秩分解矩阵）在每个Transformer块中。这大大减少了可训练参数的数量和 GPU 内存要求，因为不需要为大多数模型权重计算梯度。研究人员发现，通过关注大型语言模型的 Transformer 注意力块，使用 LoRA 进行的微调质量与完整模型微调相当，同时速度更快且需要更少的计算。

用于扩散器的 LoRA 🧨

尽管 LoRA 最初是为大型语言模型提出的，并在 transformer 块上进行了演示，但该技术也可以应用于其他地方。在稳定扩散微调的情况下，LoRA 可以应用于将图像表示与描述它们的提示相关联的交叉注意层。下图的细节（取自Stable Diffusion 论文）并不重要，只需注意黄色块是负责建立图像和文本表示之间关系的块。

据我们所知，Simo Ryu ( @cloneofsimo) 是第一个提出适用于稳定扩散的 LoRA 实现的人。请务必查看他们的 GitHub 项目以查看示例以及许多有趣的讨论和见解。

为了将 LoRA 可训练矩阵注入到与交叉注意力层一样深的模型中，过去人们需要以富有想象力（但脆弱）的方式破解扩散器的源代码。如果 Stable Diffusion 向我们展示了一件事，那就是社区总是想出办法来改变和调整模型以达到创造性目的，我们喜欢这样！出于许多其他原因，提供操纵交叉注意力层的灵活性可能是有益的，例如更容易采用xFormers等优化技术。Prompt-to-Prompt等其他创意项目可以通过一些简单的方法来访问这些层，因此我们决定为用户提供一种通用的方法。我们一直在测试从 12 月下旬开始拉取请求，并在昨天与我们的扩散器发布一起正式启动。

我们一直致力于@cloneofsimo为 Dreambooth 和全微调方法提供扩散器中的 LoRA 培训支持！这些技术提供了以下好处：

• 如前所述，训练要快得多。
• 计算要求较低。我们可以在具有 11 GB VRAM 的 2080 Ti 中创建一个完全微调的模型！
• 受过训练的重量要小得多。由于原始模型已冻结，我们注入了新层进行训练，因此我们可以将新层的权重保存为大小约为 3 MB 的单个文件。这比 UNet 模型的原始尺寸小了大约一千倍！

我们对最后一点特别兴奋。为了让用户分享他们出色的微调或梦幻模型，他们必须分享最终模型的完整副本。其他想要试用它们的用户必须在他们最喜欢的 UI 中下载经过微调的权重，这会增加大量存储和下载成本。截至今天，大约有1,000 个 Dreambooth 模型在 Dreambooth 概念库中注册，可能还有更多未在库中注册。

使用 LoRA，现在可以发布一个 3.29 MB 的文件，以允许其他人使用您的微调模型。

（h/t to @mishig25，我听到的第一个人在正常对话中将dreamboothing用作动词）。

LoRA微调

Stable Diffusion 的完整模型微调过去既缓慢又困难，这也是 Dreambooth 或 Textual Inversion 等轻量级方法变得如此流行的部分原因。使用 LoRA，在自定义数据集上微调模型要容易得多。

Diffusers 现在提供了一个LoRA 微调脚本，可以在低至 11 GB 的 GPU RAM 中运行，而无需借助 8 位优化器等技巧。这就是您如何使用它来微调使用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）上又运行了一次，这是他的最终模型，这是使用它的演示空间。

有关扩散器中 LoRA 支持的更多详细信息，请参阅我们的文档——它将始终与实施保持同步。

推理

正如我们所讨论的，LoRA 的主要优势之一是您可以通过训练比原始模型大小少几个数量级的权重来获得出色的结果。我们设计了一个推理过程，允许在未修改的稳定扩散模型权重之上加载额外的权重。让我们看看它是如何工作的。

首先，我们将使用 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 进行微调的基础模型之后，我们加载了一个正常的稳定扩散管道。我们将使用 对其进行自定义以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)

这就是神奇的地方。我们在常规模型权重之上从集线器加载 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 允许您向稳定扩散模型“教授”新概念。LoRA 与 Dreambooth 兼容，过程类似于微调，有几个优点：

• 训练速度更快。
• 我们只需要要训练的主题的几张图像（通常 5 或 10 张就足够了）。
• 如果需要，我们可以调整文本编码器，以提高对主题的保真度。

要使用 LoRA 训练 Dreambooth，您需要使用此扩散器脚本。请查看自述文件、文档和我们的超参数探索博客文章了解详细信息。

如需使用 LoRA 训练您的 Dreambooth 模型的快速、便宜且简单的方法，请查看此空间hysts。您需要复制它并分配一个 GPU 以便它运行得更快。此过程将使您不必设置自己的训练环境，您将能够在几分钟内训练您的模型！

其他方法

对轻松微调的追求并不新鲜。除了 Dreambooth 之外，文本倒置是另一种流行的方法，它试图向训练有素的稳定扩散模型教授新概念。使用 Textual Inversion 的主要原因之一是经过训练的权重也很小且易于共享。然而，它们只适用于单个主题（或一小部分主题），而 LoRA 可用于通用微调，这意味着它可以适应新的领域或数据集。

Pivotal Tuning是一种尝试将 Textual Inversion 与 LoRA 相结合的方法。首先，您使用文本反转技术向模型教授一个新概念，获得一个新的标记嵌入来表示它。然后，您使用 LoRA 训练该令牌嵌入以获得两全其美。

我们还没有使用 LoRA 探索 Pivotal Tuning。谁来迎接挑战？🤗