Diffusers 中的 LoRA 低秩适配训练:从 DreamBooth 到 Text-to-Image 的完整实战指南
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
本文围绕 🤗 Diffusers 仓库中基于 LoRA(Low-Rank Adaptation of Large Language Models)的低成本微调方案展开,系统讲解 LoRA 的核心原理、train_dreambooth_lora.py与train_text_to_image_lora.py两个训练脚本的完整用法、关键命令行参数,以及训练完成后如何用load_attn_procs/load_lora_weights加载适配器并通过scale控制融合强度。读完本文,你将能够在单张消费级 GPU 上完成 Stable Diffusion 的个性化微调,并掌握从训练到推理的完整闭环。
LoRA 是什么:用极少量可训练参数完成大模型微调
LoRA(Low-Rank Adaptation of Large Language Models) 是一种在保持内存占用较低的同时加速大模型训练的轻量级学习方法。它的核心思路是:向模型中插入一组秩分解权重矩阵对(rank-decomposition weight matrix pairs),即所谓的更新矩阵(update matrices),并且在训练过程中只训练这些新加入的权重,原始预训练权重保持冻结。
这种设计带来几个关键优势:
- 避免灾难性遗忘(catastrophic forgetting):由于预先训练好的权重始终被固定,模型在学习新数据时不会破坏原有能力;
- 参数规模极小、便于分发:秩分解矩阵的参数量远小于原始模型,训练得到的 LoRA 权重文件通常只有几百 MB,容易存储与共享,也可以与其他训练技术(如 DreamBooth)组合使用以加速训练;
- 低显存门槛:内存效率显著提升,可以在 Tesla T4、RTX 3080 或 RTX 2080 Ti 这类消费级 GPU 上直接运行微调;其中 T4 在 Kaggle 或 Google Colab 笔记本上即可免费获得。
[!WARNING] 当前版本中,LoRA 权重主要作用于模型内部的注意力(attention)层。文档明确指出:当前 LoRA 仅在
UNet2DConditionModel的注意力层中得到支持。
[!TIP] 💡 LoRA 并不局限于注意力层。原作者发现,在语言模型中仅修改注意力层即可高效获得良好性能,这正是 LoRA 权重通常被添加到模型注意力层的普遍原因。关于 LoRA 原理更深入的解释,可参考 Hugging Face 官方博客Using LoRA for effective Stable Diffusion fine-tuning。
在 Diffusers 中,cloneofsimo最早在公开的lora仓库中尝试了对 Stable Diffusion 的 LoRA 训练。随后 Diffusers 官方在 text-to-image 训练 与 DreamBooth 训练 两个示例脚本中都内置了 LoRA 支持,本文会逐一演示这两种用法。
开始之前,如果希望将模型保存到 Hub 或与社区共享,请先登录 Hugging Face 账号(没有账号可先注册):
hf auth login环境准备:安装依赖与配置 Accelerate
从源码安装 Diffusers 并安装示例依赖
训练脚本运行前,建议直接从源码安装 Diffusers,以确保使用到仓库中最新的 API:
git clone https://github.com/huggingface/diffusers cd diffusers pip install .随后进入包含训练脚本的示例目录,安装该脚本所需的依赖:
cd examples/text_to_image pip install -r requirements.txt若使用 DreamBooth LoRA 脚本,则安装 examples/dreambooth/requirements.txt 中的依赖。
初始化 Accelerate 环境
🤗 Accelerate 用于在多 GPU/TPU 或混合精度场景下辅助训练,它会根据硬件与环境自动配置训练设置。初始化方式有三种:
# 交互式配置 accelerate config# 使用默认配置(不做任何自定义选择) accelerate config default对于不支持交互式 shell 的环境(如 Jupyter Notebook),可以在代码中直接写入基础配置:
from accelerate.utils import write_basic_config write_basic_config()准备自己的训练数据集
如果想用自己的数据训练,可以参考 创建训练数据集指南,学习如何构造与训练脚本兼容的数据集(脚本同时支持从 Hub 加载数据集与从本地imagefolder目录读取数据两种方式,见 train_text_to_image_lora.py)。
训练脚本参数:LoRA 相关的两个核心参数
训练脚本通过parse_args()暴露了大量可定制参数,大多数参数都提供了表现良好的默认值。基础参数的详细说明可参考 Text-to-image 训练指南,这里重点介绍与 LoRA 直接相关的两个参数:
--rank:低秩矩阵的内部维度(inner dimension)。rank越大,可训练参数越多,模型表达能力越强,但显存与存储开销也随之上升;在 train_text_to_image_lora.py 中其默认值为4;--learning_rate:默认学习率为1e-4。得益于 LoRA 只训练少量新参数,通常可以采用比全量微调更高的学习率。
例如,要增加训练轮数,可以这样指定:
accelerate launch train_text_to_image_lora.py \ --num_train_epochs=150 \训练脚本剖析:LoRA 适配器如何接入模型
训练脚本的数据预处理与训练循环都集中在main()函数中(见 train_text_to_image_lora.py),如果要改造脚本,这里就是修改的切入点。与全量微调脚本train_text_to_image.py相比,LoRA 版本的主要差异集中在适配器的创建与过滤上。
UNet 场景:用 PEFT 的 LoraConfig 注入适配器
Diffusers 使用 PEFT 库中的LoraConfig来配置 LoRA 适配器的参数,包括rank、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())从源码可见,训练前会先将 UNet、VAE、文本编码器全部冻结(requires_grad_(False),见 train_text_to_image_lora.py),再通过add_adapter注入 LoRA 层,从而保证只有 LoRA 权重参与更新。在 fp16 混合精度训练时,还会调用cast_training_params(unet, dtype=torch.float32)将可训练参数上转为 float32,以保障数值稳定性(见 train_text_to_image_lora.py)。
target_modules中的to_q、to_k、to_v、to_out.0分别对应注意力机制中的 query、key、value 投影与输出投影,这正是 LoRA 作用于注意力层的直接体现。
文本编码器场景:SDXL 等模型的扩展支持
当需要微调文本编码器时(例如训练 Stable Diffusion XL),Diffusers 同样通过 PEFT 的LoraConfig为文本编码器添加适配器,并只过滤训练 LoRA 层:
text_lora_config = LoraConfig( r=args.rank, lora_alpha=args.rank, init_lora_weights="gaussian", target_modules=["q_proj", "k_proj", "v_proj", "out_proj"], ) text_encoder_one.add_adapter(text_lora_config) text_encoder_two.add_adapter(text_lora_config) text_lora_parameters_one = list(filter(lambda p: p.requires_grad, text_encoder_one.parameters())) text_lora_parameters_two = list(filter(lambda p: p.requires_grad, text_encoder_two.parameters()))优化器只优化 LoRA 层
优化器使用前面过滤出的lora_layers初始化,因为这些是唯一需要优化的权重(见 train_text_to_image_lora.py):
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, )默认使用torch.optim.AdamW;若安装了bitsandbytes并传入--use_8bit_adam,则会切换为bnb.optim.AdamW8bit(8 位优化器)以进一步降低显存占用(见 train_text_to_image_lora.py)。
此外,脚本通过accelerator.register_save_state_pre_hook/register_load_state_pre_hook注册了 LoRA 权重的专用保存与加载钩子(见 train_text_to_image_lora.py):保存时调用StableDiffusionPipeline.save_lora_weights写出pytorch_lora_weights.safetensors,断点续训时则通过lora_state_dict与convert_unet_state_dict_to_peft将权重还原回 PEFT 适配器。除 LoRA 层的设置外,训练脚本整体流程与train_text_to_image.py基本一致。
实战一:Text-to-Image 的 LoRA 微调
启动训练
下面用 Naruto BLIP captions 数据集微调stable-diffusion-v1-5/stable-diffusion-v1-5,训练一个能生成火影忍者风格角色的模型。
首先设置环境变量:MODEL_NAME指定基础模型,DATASET_NAME指定数据集,OUTPUT_DIR与HUB_MODEL_ID可选,分别指定模型在本地与 Hub 上的保存位置:
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"训练完成后,脚本会在输出目录中生成以下文件:
- 模型检查点(checkpoint)
pytorch_lora_weights.safetensors(训练得到的 LoRA 权重)
多卡训练时,请在accelerate launch命令中追加--multi_gpu参数。
[!WARNING] 在 11GB 显存的 2080 Ti GPU 上,完整训练一次大约需要 5 小时。
启动训练:
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推理:加载 LoRA 权重生成图像
训练完成后,用AutoPipelineForText2Image加载基础模型与 LoRA 权重进行推理:
from diffusers import AutoPipelineForText2Image import torch pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda") # 或 "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]其中load_lora_weights是 Diffusers 在 lora_pipeline.py 中提供的通用 LoRA 加载接口,支持从本地路径或 Hub 仓库加载pytorch_lora_weights.safetensors。
实战二:DreamBooth + LoRA 微调
准备环境变量
DreamBooth 是另一种常见的个性化微调方案,通过少量(通常 3~5 张)同一主体的图片配合唯一标识符(如sks dog)来学习特定对象。与 LoRA 结合后,训练速度与显存开销都大幅下降。
以下基于 train_dreambooth_lora.py 演示。训练前需要设置MODEL_NAME(基础模型)与INSTANCE_DIR(实例图片目录,即待学习主体所在的本地文件夹)。OUTPUT_DIR与HUB_MODEL_ID可选,分别指定模型在本地与 Hub 上的保存位置:
export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5" export INSTANCE_DIR="./dog" # 包含 sks dog 示例图片的本地目录 export OUTPUT_DIR="/sddata/finetune/lora/dog" export HUB_MODEL_ID="dog-lora"关键训练标志说明
开始训练前,需要了解几个重要参数:
--push_to_hub:将训练好的 LoRA 权重保存到 Hub;--report_to=wandb:将训练结果与日志上报到 Weights & Biases 仪表盘;--learning_rate=1e-04:LoRA 训练通常可以采用比全量微调更高的学习率;--rank:低秩矩阵维度,默认4(见 train_dreambooth_lora.py);--lora_dropout:LoRA 层的 dropout 概率,默认0.0(见 train_dreambooth_lora.py)。
启动训练
accelerate launch train_dreambooth_lora.py \ --pretrained_model_name_or_path=$MODEL_NAME \ --instance_data_dir=$INSTANCE_DIR \ --output_dir=$OUTPUT_DIR \ --instance_prompt="a photo of sks dog" \ --resolution=512 \ --train_batch_size=1 \ --gradient_accumulation_steps=1 \ --checkpointing_steps=100 \ --learning_rate=1e-4 \ --report_to="wandb" \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --max_train_steps=500 \ --validation_prompt="A photo of sks dog in a bucket" \ --validation_epochs=50 \ --seed="0" \ --push_to_hub推理:加载注意力处理器并控制融合比例
训练完成后,通过StableDiffusionPipeline加载基础模型,然后在基础模型权重之上叠加 DreamBooth 微调得到的 LoRA 权重,并将 pipeline 迁移到 GPU 以加速推理:
import torch from diffusers import StableDiffusionPipeline model_base = "stable-diffusion-v1-5/stable-diffusion-v1-5" pipe = StableDiffusionPipeline.from_pretrained(model_base, dtype=torch.float16) pipe.unet.load_attn_procs(model_path) pipe.to("cuda")[!TIP] 💡
scale参数用于控制 LoRA 权重与冻结的预训练权重合并时的比例:scale=0表示完全不使用 LoRA 权重(等价于只用原始模型),scale=1表示完全使用微调后的 LoRA 权重,0 到 1 之间的值则是在两者结果之间进行插值。
# 使用 50% 的 LoRA 权重 + 50% 的基础模型权重 image = pipe( "A picture of a sks dog in a bucket.", num_inference_steps=25, guidance_scale=7.5, cross_attention_kwargs={"scale": 0.5}, ).images[0] # 完全使用微调后的 LoRA 权重 image = pipe("A picture of a sks dog in a bucket.", num_inference_steps=25, guidance_scale=7.5).images[0] image.save("bucket-dog.png")load_attn_procs 背后的实现
load_attn_procs定义于 src/diffusers/loaders/unet.py,用于将预训练的注意力处理器(attention processor)层加载进UNet2DConditionModel。其入参可以是 Hub 上的模型 ID、包含权重的本地目录路径,或直接传入 torch state dict;同时支持weight_name指定序列化的权重文件名、cache_dir指定缓存目录、revision指定模型版本等参数。它的执行流程是:优先尝试加载.safetensors权重,再通过set_peft_model_state_dict将状态字典还原到 PEFT 适配器结构中(参见 train_text_to_image_lora.py 中的加载钩子实现)。
Diffusers 还提供更高级的fuse_lora方法(见 src/diffusers/loaders/peft.py),可以在推理前将 LoRA 权重**融合(fuse)**进原始权重中,从而在不额外调用 LoRA 前向逻辑的情况下获得与scale参数等价的控制效果——传入lora_scale即对应这里的scale语义。
下一步学习
恭喜你完成了 LoRA 微调的全流程!想进一步深入,可以参考以下主题:
- 学习如何加载 Kohya、TheLastBen 等社区训练器产出的不同 LoRA 格式;
- 学习如何使用 PEFT 组合多个 LoRA 进行推理,实现风格叠加;
- 参考 DreamBooth 训练指南 与 Text-to-image 训练指南,对比 LoRA 方案与全量微调方案在参数、脚本结构与训练流程上的异同。
【免费下载链接】diffusers🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考