跳到正文
原文
Hugging Face Blog·· 2023-01-26精选AI 评分82

用 LoRA 高效微调 Stable Diffusion：diffusers 支持与实测流程

Using LoRA for Efficient Stable Diffusion Fine-Tuning

AI 导读

Hugging Face 在 diffusers 中正式加入 LoRA 微调支持，可在 11 GB 显存的 2080 Ti 上训练 Stable Diffusion，微调权重文件仅约 3 MB，比 UNet 原模型小约一千倍。

推荐理由

以 11 GB 显存训 Stable Diffusion 的完整命令与推理代码为例，展示 LoRA 把微调权重压到 3 MB 的可行路径。

正文 · AI 翻译

LoRA：大型语言模型的低秩适配是微软研究人员提出的一项新技术，旨在解决大型语言模型微调的问题。像 GPT-3 这样拥有数十亿参数的强大模型，为了适应特定任务或领域而进行微调的成本高得令人望而却步。LoRA 提出冻结预训练模型的权重，并在每个 transformer 块中注入可训练层（秩分解矩阵）。这大大减少了可训练参数的数量和 GPU 内存需求，因为大多数模型权重不需要计算梯度。研究人员发现，通过聚焦于大型语言模型的 Transformer 注意力块，LoRA 的微调质量与完整模型微调相当，同时速度更快、所需计算量更少。

适用于 Diffusers 的 LoRA 🧨

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

Latent Diffusion Architecture

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

为了将 LoRA 可训练矩阵注入到模型中如交叉注意力层那样深的位置，人们过去需要以富有想象力（但脆弱）的方式修改 diffusers 的源代码。如果 Stable Diffusion 向我们展示了一件事，那就是社区总能想出办法来弯曲和调整模型以实现创意目的，我们很喜欢这一点！提供操作交叉注意力层的灵活性可能出于许多其他原因而有益，例如更容易采用 xFormers 等优化技术。其他创意项目如 Prompt-to-Prompt 也可以通过某种简便方式访问这些层，因此我们决定为用户提供一种通用的方法来实现这一点。我们从 12 月下旬开始测试该拉取请求，它已随我们昨天的 diffusers 发布正式推出。

我们一直在与 @cloneofsimo 合作，在 diffusers 中为 Dreambooth 和完整微调方法提供 LoRA 训练支持！这些技术提供了以下好处：

  • 如前所述，训练速度要快得多。
  • 计算需求更低。我们可以在拥有 11 GB 显存的 2080 Ti 上创建一个完整微调模型！
  • 训练后的权重小得多得多。因为原始模型被冻结，我们注入新的层进行训练，我们可以将新层的权重保存为一个大小约为 3 MB 的单个文件。这大约比 UNet 模型的原始大小小一千倍！

我们对最后一点尤为兴奋。为了让用户分享他们出色的微调或 dreamboothed 模型，他们不得不分享最终模型的完整副本。其他想要尝试这些模型的用户则需在各自偏好的界面中下载微调后的权重，这导致了存储和下载成本的累加。截至目前，大约有 1,000 个 Dreambooth 模型在 Dreambooth Concepts Library 注册，可能还有更多未注册的模型。

借助 LoRA，现在只需发布一个 3.29 MB 的文件，就能让他人使用你的微调模型。

（感谢 @mishig25，我首次在日常对话中听到有人将 dreamboothing 作为动词使用。）

LoRA 微调

过去，对 Stable Diffusion 进行全模型微调既慢又难，这也是诸如 Dreambooth 或文本反转等更轻量级方法变得如此受欢迎的部分原因。有了 LoRA，基于自定义数据集微调模型变得容易得多。

Diffusers 现在提供了一个LoRA 微调脚本，即使在低至 11 GB 的 GPU 内存下也能运行，而无需使用诸如 8 位优化器等技巧。以下是使用Lambda Labs 宝可梦数据集微调模型的方法：

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 内存）上耗时约 5 小时。我没有尝试优化超参数，欢迎大家自行尝试！Sayak 在 T4（16 GB 内存）上进行了另一次运行，这是他的最终模型，还有一个使用该模型的演示空间。

Sample outputs from Sayak's LoRA model

关于 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 进行 Dreamboothing

Dreambooth 允许你向 Stable Diffusion 模型“传授”新概念。LoRA 与 Dreambooth 兼容，其流程与微调类似，但有几个优势：

  • 训练速度更快。
  • 我们只需要少量想要训练的主题图像（通常 5 到 10 张就足够了）。
  • 如果需要，我们还可以微调文本编码器，以增强对主题的保真度。

要使用 LoRA 训练 Dreambooth，你需要使用 这个 diffusers 脚本。详情请参阅 README、文档 以及 我们的超参数探索博客文章。

如需一种快速、经济且简便的方法来使用 LoRA 训练你的 Dreambooth 模型，请查看 hysts 制作的 这个 Space。你需要复制它并分配一个 GPU 以便快速运行。此流程将省去你自行搭建训练环境的麻烦，让你能在几分钟内完成模型训练！

其他方法

对简便微调的探索并非新鲜事。除了 Dreambooth 之外，文本反转 是另一种流行的方法，它试图向已训练好的 Stable Diffusion 模型传授新概念。使用文本反转的主要原因之一是训练出的权重也很小且易于分享。然而，它们仅适用于单个主题（或少数几个主题），而 LoRA 可用于通用微调，这意味着它可以适应新的领域或数据集。

Pivotal Tuning 是一种尝试将文本反转与 LoRA 结合的方法。首先，你使用文本反转技术向模型传授一个新概念，获得一个新的 token 嵌入来表示它。然后，你使用 LoRA 训练该 token 嵌入，从而兼得两者之长。

我们尚未探索使用 LoRA 的 Pivotal Tuning。谁愿意接受这个挑战呢？🤗

来源：Hugging Face Blog · hf-mirror.com

#教程/实践#图像生成#开源生态#Hugging Face#Microsoft