Diffusers documentation

LoRA

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

LoRA

此功能目前处于实验阶段,API 未来可能会发生变化。

LoRA(大型语言模型的低秩适配)是一种流行且轻量级的训练技术,可大幅减少可训练参数的数量。其原理是在模型中插入数量较少的新型权重,并且仅对这些权重进行训练。这使得使用 LoRA 进行训练的速度更快、内存效率更高,并生成更小的模型权重(仅几百 MB),从而更易于存储和共享。LoRA 还可以与 DreamBooth 等其他训练技术结合使用,以加快训练速度。

LoRA 非常通用,并支持 DreamBooth、Kandinsky 2.2、Stable Diffusion XL 以及 text-to-image。

This guide will explore the train_text_to_image_lora.py script to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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

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

cd examples/text_to_image
pip install -r requirements.txt

🤗 Accelerate 是一个用于在多个 GPU/TPU 上训练或使用混合精度训练的库。它会根据你的硬件和环境自动配置训练设置。请查看 🤗 Accelerate 的 快速入门 以了解更多信息。

初始化一个 🤗 Accelerate 环境:

accelerate config

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

accelerate config default

或者,如果你的环境不支持交互式 shell(例如 notebook),你可以使用:

from accelerate.utils import write_basic_config

write_basic_config()

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

以下各部分重点介绍了训练脚本中对于理解如何修改脚本至关重要的内容,但并未详细涵盖脚本的每个方面。如果您希望了解更多,请随时阅读脚本,并告诉我们您是否有任何疑问或顾虑。

脚本参数

训练脚本包含许多参数,可帮助你自定义训练过程。所有参数及其说明均位于 parse_args() 函数中。大多数参数都提供了默认值,效果通常不错,但你也可以在训练命令中自行设置所需的值。

例如,要增加训练轮数:

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

许多基础且重要的参数已在文本到图像训练指南中介绍,因此本指南仅聚焦于与 LoRA 相关的参数:

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

训练脚本

数据集预处理代码和训练循环位于 main() 函数中,如果需要调整训练脚本,请在此处进行修改。

与脚本参数类似,文本到图像训练指南中提供了训练脚本的逐步讲解。本指南则着重介绍脚本中与 LoRA 相关的部分。

UNet
文本编码器

Diffusers 使用 PEFT 库中的 LoraConfig 来配置 LoRA 适配器的参数,例如秩(rank)、alpha 以及将 LoRA 权重插入哪些模块。该适配器会被添加到 UNet 中,并且只有 LoRA 层会在 lora_layers 中被筛选出来以进行优化。

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())

optimizer 使用 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,
)

Aside from setting up the LoRA layers, the training script is more or less the same as train_text_to_image.py!

运行脚本

完成所有更改或对默认配置满意后,即可启动训练脚本!🚀

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

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

如果要在多块 GPU 上进行训练,请在 accelerate launch 命令中添加 --multi_gpu 参数。

在配备 11GB 显存的 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", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_lora_weights("path/to/lora/model", weight_name="pytorch_lora_weights.safetensors")
image = pipeline("A naruto with blue eyes").images[0]

后续步骤

恭喜您使用 LoRA 训练了一个新模型!如需了解如何使用新模型,以下指南可能对您有所帮助:

  • 了解如何加载使用 Kohya 和 TheLastBen 等社区训练器训练的不同 LoRA 格式。
  • 了解如何使用 PEFT 进行推理,并组合多个 LoRA。
Update on GitHub