Diffusers 文档

LoRA

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始

LoRA

这是实验性的,API 未来可能会更改。

LoRA(大型语言模型的低秩自适应)是一种流行的轻量级训练技术,可显著减少可训练参数的数量。它的工作原理是在模型中插入少量新权重,并且仅训练这些权重。这使得使用 LoRA 进行训练更快、更节省内存,并产生更小的模型权重(几百 MB),这些权重更容易存储和共享。LoRA 还可以与其他训练技术(如 DreamBooth)结合使用以加速训练。

LoRA 用途广泛,并支持 DreamBoothKandinsky 2.2Stable Diffusion XLtext-to-imageWuerstchen

本指南将探讨 train_text_to_image_lora.py 脚本,以帮助您更熟悉它,以及如何将其调整为自己的用例。

在运行脚本之前,请确保从源代码安装库

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

导航到包含训练脚本的示例文件夹,并为您正在使用的脚本安装所需的依赖项

PyTorch
Flax
cd examples/text_to_image
pip install -r requirements.txt

🤗 Accelerate 是一个库,可帮助您在多个 GPU/TPU 上或使用混合精度进行训练。它将根据您的硬件和环境自动配置您的训练设置。请查看 🤗 Accelerate 快速导览以了解更多信息。

初始化一个 🤗 Accelerate 环境

accelerate config

要设置一个默认的 🤗 Accelerate 环境,而无需选择任何配置

accelerate config default

或者,如果您的环境不支持交互式 shell(如笔记本),则可以使用

from accelerate.utils import write_basic_config

write_basic_config()

最后,如果您想在自己的数据集上训练模型,请查看创建用于训练的数据集指南,以了解如何创建与训练脚本一起使用的数据集。

以下部分重点介绍了训练脚本中对于理解如何修改脚本很重要的部分,但它并未详细介绍脚本的每个方面。如果您有兴趣了解更多信息,请随时阅读脚本,并告知我们您是否有任何问题或疑虑。

脚本参数

训练脚本有许多参数可以帮助您自定义训练运行。所有参数及其描述都可以在 parse_args() 函数中找到。大多数参数都提供了效果不错的默认值,但如果您愿意,也可以在训练命令中设置自己的值。

例如,要增加要训练的 epoch 数量

accelerate launch train_text_to_image_lora.py \
  --num_train_epochs=150 \

许多基本且重要的参数在文本到图像训练指南中进行了描述,因此本指南仅关注 LoRA 相关参数。

  • --rank:要训练的低秩矩阵的内部维度;更高的秩意味着更多的可训练参数
  • --learning_rate:默认学习率是 1e-4,但使用 LoRA,您可以使用更高的学习率

训练脚本

数据集预处理代码和训练循环在 main() 函数中找到,如果您需要调整训练脚本,则可以在此处进行更改。

与脚本参数一样,文本到图像训练指南中提供了训练脚本的演练。相反,本指南着眼于脚本中与 LoRA 相关的部分。

UNet
文本编码器

Diffusers 使用 PEFT 库中的 ~peft.LoraConfig 来设置 LoRA 适配器的参数,例如秩、alpha 以及将 LoRA 权重插入到哪些模块中。适配器被添加到 UNet,并且仅在 lora_layers 中过滤 LoRA 层以进行优化。

unet_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)

unet.add_adapter(unet_lora_config)
lora_layers = filter(lambda p: p.requires_grad, unet.parameters())

优化器使用 lora_layers 初始化,因为这些是唯一要优化的权重

optimizer = optimizer_cls(
    lora_layers,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

除了设置 LoRA 层之外,训练脚本或多或少与 train_text_to_image.py 相同!

启动脚本

一旦您完成所有更改或对默认配置感到满意,您就可以启动训练脚本了!🚀

让我们在 Naruto BLIP captions 数据集上进行训练,以生成您自己的火影忍者角色。将环境变量 MODEL_NAMEDATASET_NAME 分别设置为模型和数据集。您还应该在 OUTPUT_DIR 中指定模型的保存位置,以及使用 HUB_MODEL_ID 在 Hub 上保存模型的名称。该脚本会创建以下文件并将其保存到您的存储库:

  • 保存的模型检查点
  • pytorch_lora_weights.safetensors (训练后的 LoRA 权重)

如果您在多个 GPU 上进行训练,请将 --multi_gpu 参数添加到 accelerate launch 命令。

在配备 11GB VRAM 的 2080 Ti GPU 上,完整的训练运行需要约 5 小时。

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/naruto"
export HUB_MODEL_ID="naruto-lora"
export DATASET_NAME="lambdalabs/naruto-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="A naruto with blue eyes." \
  --seed=1337

训练完成后,您可以使用您的模型进行推理。

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipeline.load_lora_weights("path/to/lora/model", weight_name="pytorch_lora_weights.safetensors")
image = pipeline("A naruto with blue eyes").images[0]

后续步骤

祝贺您使用 LoRA 训练了一个新模型!要了解有关如何使用新模型的更多信息,以下指南可能会有所帮助:

< > 在 GitHub 上更新