ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Diffusers 中的 DreamBooth 微调指南:从基础训练到显存优化

Diffusers 中的 DreamBooth 微调指南:从基础训练到显存优化 Diffusers 中的 DreamBooth 微调指南从基础训练到显存优化【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers导读DreamBooth 是一种只需少量示例图像通常 3~5 张即可对 Stable Diffusion 这类 text-to-image 模型进行个性化微调的训练技术。它通过在提示词中绑定一个特殊标识词例如sks dog与示例图像使模型能够在各种场景、姿态和视角下生成该主题的上下文相关图像。本文将基于 Hugging Face Diffusers 仓库中的韩文官方文档docs/source/ko/training/dreambooth.md结合仓库内完整的训练脚本源码系统地讲解 DreamBooth 的核心原理、环境搭建、完整训练命令、检查点机制以及从 16GB 到 8GB 显存的分级优化方案让你能够直接复现并应用到自己的个性化图像生成任务中。1. DreamBooth 基本原理与项目支持DreamBooth 由论文 DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation 提出其核心思想是利用极少量3~5 张某一主题的图像通过微调整个扩散模型让模型学会将提示词中的特殊标识符与示例图像关联起来。微调完成后模型即可在多样化的场景、姿态与视角下生成该主题的图像。在 Diffusers 仓库中DreamBooth 训练功能由 examples/dreambooth/ 目录下的脚本提供核心文件包括examples/dreambooth/train_dreambooth.py主训练脚本支持 UNet 单独微调、文本编码器联合微调、先验保留损失prior preservation loss、LoRA、检查点保存与恢复等功能examples/dreambooth/requirements.txt训练依赖清单examples/dreambooth/test_dreambooth.py基于 tiny 模型的冒烟测试验证脚本可运行性examples/dreambooth/README.md仓库内的英文训练说明。本指南以CompVis/stable-diffusion-v1-4为基底模型演示不同 GPU 显存规格下的微调方案。2. 环境准备与依赖安装2.1 安装 Diffusers 与训练依赖运行训练脚本前需要先安装库的训练依赖官方建议直接从mainGitHub 分支安装 Diffuserspip install githttps://github.com/huggingface/diffusers pip install -U -r diffusers/examples/dreambooth/requirements.txt其中requirements.txt的具体内容见 examples/dreambooth/requirements.txt如下accelerate0.16.0 torchvision transformers4.25.1 ftfy tensorboard Jinja2 peft0.7.0要点说明accelerate0.16.0用于多卡/混合精度训练同时也是新版检查点推理方案的版本分界线transformers提供 CLIP 文本编码器与 tokenizerpeft0.7.0用于 LoRA 相关训练路径tensorboard作为默认的日志记录后端。2.2 安装 xFormers可选但推荐xFormers 不是训练的硬性要求但官方建议尽可能安装见 docs/source/ko/optimization/xformers.md因为其 memory-efficient attention 机制能够加快训练速度并降低显存占用pip install xformers[!NOTE] 根据 docs/source/ko/optimization/xformers.md 中的记录xFormersv0.0.16在 GPU 上进行训练微调或 DreamBooth时可能存在已知问题如果遇到问题请升级到更高版本至少 0.0.17或安装 development 版本。2.3 初始化 Accelerate 环境训练脚本通过 Accelerate 管理分布式训练与混合精度。使用交互式方式初始化accelerate config若不希望做任何配置直接使用默认环境accelerate config default如果当前环境不支持交互式 shell例如 notebook可以通过 Python API 写入基础配置from accelerate.utils import write_basic_config write_basic_config()3. 基础微调仅训练 UNet[!WARNING] DreamBooth 微调对超参数非常敏感且容易过拟合。建议参考官方 Training Stable Diffusion with Dreambooth using Diffusers 博客中的推荐设置来选择合适超参数。3.1 准备训练数据以几张狗狗图像为例可自行收集 3~5 张同一主题的图像下载后放入目录并通过环境变量指定路径export MODEL_NAMECompVis/stable-diffusion-v1-4 export INSTANCE_DIRpath_to_training_images export OUTPUT_DIRpath_to_saved_modelMODEL_NAME基底模型名称Hub 模型 ID或本地路径INSTANCE_DIR训练实例图像的目录OUTPUT_DIR模型保存目录。3.2 启动训练accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path$MODEL_NAME \ --instance_data_dir$INSTANCE_DIR \ --output_dir$OUTPUT_DIR \ --instance_prompta photo of sks dog \ --resolution512 \ --train_batch_size1 \ --gradient_accumulation_steps1 \ --learning_rate5e-6 \ --lr_schedulerconstant \ --lr_warmup_steps0 \ --max_train_steps400关键参数解析对应 examples/dreambooth/train_dreambooth.py 中parse_args()的定义参数默认值说明--pretrained_model_name_or_path必填Hub 上的模型 ID 或本地预训练模型路径--instance_data_dir必填包含实例训练图像的文件夹--instance_prompt必填包含特殊标识词的文本提示词如a photo of sks dog--output_dirdreambooth-model模型与检查点的输出目录--resolution512输入图像分辨率训练/验证数据集中的所有图像都会被缩放至该尺寸--train_batch_size4每个设备的训练 batch size--gradient_accumulation_steps1梯度累积步数累积后再执行 backward/update--learning_rate5e-6初始学习率在 warmup 之后--lr_schedulerconstant学习率调度器可选linear、cosine、cosine_with_restarts、polynomial、constant、constant_with_warmup--lr_warmup_steps500学习率 warmup 步数脚本默认值本命令显式设为 0--max_train_stepsNone总训练步数指定后覆盖num_train_epochs[!NOTE] 训练脚本内部会加载 VAE、文本编码器与 UNet其中 VAE 固定vae.requires_grad_(False)未开启--train_text_encoder时文本编码器也被冻结text_encoder.requires_grad_(False)仅 UNet 参与训练单卡 batch size 为 1、仅 400 步的配置非常适合小显存环境入门。4. 先验保留损失Prior Preservation Loss微调4.1 原理为了防止过拟合与language drift语言漂移DreamBooth 引入了先验保留机制在训练过程中混入同一类别的其他图像class images。巧妙之处在于这些类别图像可以直接用 Stable Diffusion 模型本身生成训练脚本会自动将生成的类别图像保存到指定的本地路径。依据论文建议先验保留需要生成num_epochs * num_samples张类别图像实践中200~300 张通常效果不错。4.2 训练命令export MODEL_NAMECompVis/stable-diffusion-v1-4 export INSTANCE_DIRpath_to_training_images export CLASS_DIRpath_to_class_images export OUTPUT_DIRpath_to_saved_model accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path$MODEL_NAME \ --instance_data_dir$INSTANCE_DIR \ --class_data_dir$CLASS_DIR \ --output_dir$OUTPUT_DIR \ --with_prior_preservation --prior_loss_weight1.0 \ --instance_prompta photo of sks dog \ --class_prompta photo of dog \ --resolution512 \ --train_batch_size1 \ --gradient_accumulation_steps1 \ --learning_rate5e-6 \ --lr_schedulerconstant \ --lr_warmup_steps0 \ --num_class_images200 \ --max_train_steps800相关参数说明--with_prior_preservation开启先验保留损失store_true 开关--class_data_dir类别图像目录。开启先验保留后必填若目录中图像不足--num_class_images脚本会自动生成补充--class_prompt描述类别图像的提示词开启先验保留后必填--prior_loss_weight默认1.0控制先验保留损失对模型的影响权重--num_class_images默认100先验保留所需的最少类别图像数。源码中会在cur_class_images args.num_class_images时使用PromptDataset按class_prompt批量采样生成缺失图像见 examples/dreambooth/train_dreambooth.py--sample_batch_size默认4生成类别图像时的采样 batch size。4.3 源码级实现细节从源码看先验保留损失的实现位于训练循环中examples/dreambooth/train_dreambooth.py类别图像与实例图像通过collate_fn拼接在同一 batch 中避免两次前向传播见 collate_fn对model_pred与target按 batch 维度切分分别计算实例损失与先验损失最终损失为loss loss args.prior_loss_weight * prior_loss。此外训练脚本还支持其他高级参数--snr_gammaMin-SNR 加权推荐值 5.0、--offset_noise、--center_crop、--scale_lr、--use_8bit_adam等均可按需在命令行中追加。5. 联合微调文本编码器与 UNet实验表明详见官方博客 Training Stable Diffusion with Dreambooth using Diffusers联合微调文本编码器能显著提升生成质量尤其是在人脸图像生成场景下。[!WARNING] 训练文本编码器需要额外显存16GB GPU 无法运行使用该选项至少需要24GB VRAM。通过向训练脚本传递--train_text_encoder参数即可同时微调文本编码器与 UNetexport MODEL_NAMECompVis/stable-diffusion-v1-4 export INSTANCE_DIRpath_to_training_images export CLASS_DIRpath_to_class_images export OUTPUT_DIRpath_to_saved_model accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path$MODEL_NAME \ --train_text_encoder \ --instance_data_dir$INSTANCE_DIR \ --class_data_dir$CLASS_DIR \ --output_dir$OUTPUT_DIR \ --with_prior_preservation --prior_loss_weight1.0 \ --instance_prompta photo of sks dog \ --class_prompta photo of dog \ --resolution512 \ --train_batch_size1 \ --use_8bit_adam \ --gradient_checkpointing \ --learning_rate2e-6 \ --lr_schedulerconstant \ --lr_warmup_steps0 \ --num_class_images200 \ --max_train_steps800此配置下的额外要点联合训练文本编码器时学习率通常需要调低示例中使用2e-6--use_8bit_adam与--gradient_checkpointing用于缓解显存压力从源码看examples/dreambooth/train_dreambooth.py开启--train_text_encoder后优化器参数集合为itertools.chain(unet.parameters(), text_encoder.parameters())两个模型同时参与更新文本编码器训练时必须保持 float32 精度源码对此有明确校验见 examples/dreambooth/train_dreambooth.py。6. LoRA 微调对于大模型还可以使用LoRALow-Rank Adaptation of Large Language Models这一加速微调技术来降低训练成本。Diffusers 在 examples/dreambooth/ 目录下提供了train_dreambooth_lora.py脚本对应 docs/source/ko/training/lora.md 中 DreamBooth 训练 一节其基本用法为accelerate launch train_dreambooth_lora.py \ --pretrained_model_name_or_path$MODEL_NAME \ --instance_data_dir$INSTANCE_DIR \ --output_dir$OUTPUT_DIR \ --instance_prompta photo of sks dog \ --resolution512 \ --train_batch_size1 \ --gradient_accumulation_steps1 \ --max_train_steps800 \ --learning_rate1e-4 \ --lr_schedulerconstant \ --lr_warmup_steps0LoRA 训练只更新低秩适配矩阵显存占用与训练时长远小于全量微调训练产物为轻量 LoRA 权重适合分享与复用。详细说明请参考 docs/source/ko/training/lora.md。7. 训练中的检查点保存与恢复DreamBooth 训练过程中极易过拟合因此周期性保存中间检查点非常有用——某个中间检查点可能比最终模型效果更好7.1 启用检查点保存在训练命令中加入--checkpointing_steps500该参数默认值即为500见 examples/dreambooth/train_dreambooth.py表示每 500 步保存一次完整训练状态。检查点保存在output_dir的子目录中目录名以checkpoint-前缀加已完成的步数命名例如checkpoint-1500表示第 1500 个训练步后保存的检查点。补充说明可通过--checkpoints_total_limit限制最多保留的检查点数量超出时脚本会按步数从小到大自动删除最早的检查点见 examples/dreambooth/train_dreambooth.py检查点内容不仅包含模型权重还包含优化器、数据加载器以及学习率调度器的状态通过accelerator.save_state保存。7.2 从检查点恢复训练使用--resume_from_checkpoint参数并指定检查点名称即可恢复训练--resume_from_checkpointcheckpoint-1500也可以使用特殊字符串latest自动从步数最多的最新检查点恢复。源码逻辑会扫描output_dir下所有checkpoint-*目录并按其步数排序选取见 examples/dreambooth/train_dreambooth.py。恢复后可根据需要调整部分超参数。7.3 用保存的检查点进行推理保存的检查点格式专为训练恢复设计包含模型权重以及优化器、数据加载器、学习率的状态。当安装的accelerate0.16.0时可以直接加载中间检查点中的模型组件进行推理from diffusers import DiffusionPipeline, UNet2DConditionModel from transformers import CLIPTextModel import torch # 使用与训练时相同的参数model, revision加载管道 model_id CompVis/stable-diffusion-v1-4 unet UNet2DConditionModel.from_pretrained(/sddata/dreambooth/daruma-v2-1/checkpoint-100/unet) # 如果训练时使用了 --train_text_encoder请务必加载文本编码器 text_encoder CLIPTextModel.from_pretrained(/sddata/dreambooth/daruma-v2-1/checkpoint-100/text_encoder) pipeline DiffusionPipeline.from_pretrained(model_id, unetunet, text_encodertext_encoder, dtypetorch.float16) pipeline.to(cuda) # 执行推理、保存或推送到 Hub pipeline.save_pretrained(dreambooth-pipeline)当安装的accelerate0.16.0时需要先将检查点转换为推理管道from accelerate import Accelerator from diffusers import DiffusionPipeline # 使用与训练时相同的参数model, revision加载管道 model_id CompVis/stable-diffusion-v1-4 pipeline DiffusionPipeline.from_pretrained(model_id) accelerator Accelerator() # 如果初始训练使用了 --train_text_encoder则传入文本编码器 unet, text_encoder accelerator.prepare(pipeline.unet, pipeline.text_encoder) # 从检查点路径恢复状态这里必须使用绝对路径 accelerator.load_state(/sddata/dreambooth/daruma-v2-1/checkpoint-100) # 用 unwrapped 模型重建管道直接赋值给 .unet 和 .text_encoder 也应可行 pipeline DiffusionPipeline.from_pretrained( model_id, unetaccelerator.unwrap_model(unet), text_encoderaccelerator.unwrap_model(text_encoder), ) # 执行推理、保存或推送到 Hub pipeline.save_pretrained(dreambooth-pipeline)从源码角度看检查点之所以能直接以子目录形式保存unet/与text_encoder/是因为脚本注册了自定义的save_model_hook/load_model_hook见 examples/dreambooth/train_dreambooth.py使accelerator.save_state能以 Diffusers/Transformers 的标准save_pretrained格式序列化模型。8. 不同 GPU 显存规格的优化方案根据硬件条件可以将 DreamBooth 优化到16GB 甚至 8GB的 GPU 上运行。8.1 xFormers 内存高效注意力xFormers 是优化 Transformers 的工具箱其中包含 Diffusers 使用的 memory-efficient attention 机制。安装 xFormers 后在训练命令中加入--enable_xformers_memory_efficient_attention8.2 梯度置 None另一种降低显存的方式是在optimizer.zero_grad时将梯度设为None而非 0。但注意这会改变某些行为若遇到问题请移除该参数--set_grads_to_none源码对应实现为optimizer.zero_grad(set_to_noneargs.set_grads_to_none)见 examples/dreambooth/train_dreambooth.py。8.3 16GB GPU 方案借助梯度检查点gradient checkpointing与 bitsandbytes 的 8 位优化器可在 16GB GPU 上训练。先安装 bitsandbytespip install bitsandbytes然后在训练命令中指定--use_8bit_adamexport MODEL_NAMECompVis/stable-diffusion-v1-4 export INSTANCE_DIRpath_to_training_images export CLASS_DIRpath_to_class_images export OUTPUT_DIRpath_to_saved_model accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path$MODEL_NAME \ --instance_data_dir$INSTANCE_DIR \ --class_data_dir$CLASS_DIR \ --output_dir$OUTPUT_DIR \ --with_prior_preservation --prior_loss_weight1.0 \ --instance_prompta photo of sks dog \ --class_prompta photo of dog \ --resolution512 \ --train_batch_size1 \ --gradient_accumulation_steps2 --gradient_checkpointing \ --use_8bit_adam \ --learning_rate5e-6 \ --lr_schedulerconstant \ --lr_warmup_steps0 \ --num_class_images200 \ --max_train_steps800源码中--use_8bit_adam会将优化器切换为bitsandbytes.optim.AdamW8bit见 examples/dreambooth/train_dreambooth.py。8.4 12GB GPU 方案在 12GB GPU 上运行需要组合使用梯度检查点、8 位优化器、xFormers并将梯度设为Noneexport MODEL_NAMECompVis/stable-diffusion-v1-4 export INSTANCE_DIRpath-to-instance-images export CLASS_DIRpath-to-class-images export OUTPUT_DIRpath-to-save-model accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path$MODEL_NAME \ --instance_data_dir$INSTANCE_DIR \ --class_data_dir$CLASS_DIR \ --output_dir$OUTPUT_DIR \ --with_prior_preservation --prior_loss_weight1.0 \ --instance_prompta photo of sks dog \ --class_prompta photo of dog \ --resolution512 \ --train_batch_size1 \ --gradient_accumulation_steps1 --gradient_checkpointing \ --use_8bit_adam \ --enable_xformers_memory_efficient_attention \ --set_grads_to_none \ --learning_rate2e-6 \ --lr_schedulerconstant \ --lr_warmup_steps0 \ --num_class_images200 \ --max_train_steps8008.5 8GB GPU 方案DeepSpeed 张量卸载对于 8GB GPU可以使用 DeepSpeed 将部分张量从 VRAM 卸载到 CPU 或 NVME从而在更少显存下训练。首先配置 Accelerate 环境并在配置过程中选择启用 DeepSpeedaccelerate config将DeepSpeed stage 2与fp16 混合精度结合并将模型参数与优化器状态都卸载到 CPU即可在8GB VRAM 以下训练。代价是需要更多系统 RAM约 25GB。更多配置选项请参考 DeepSpeed 官方文档。同时需要将默认的 Adam 优化器替换为 DeepSpeed 优化版的deepspeed.ops.adam.DeepSpeedCPUAdam它针对 CPU 卸载场景做了大幅提速。启用DeepSpeedCPUAdam要求系统中的 CUDA toolchain 版本与 PyTorch 自带的版本一致。[!NOTE] 8 位优化器bitsandbytes目前似乎与 DeepSpeed 不兼容二者不要同时使用。8GB 场景的完整训练命令export MODEL_NAMECompVis/stable-diffusion-v1-4 export INSTANCE_DIRpath_to_training_images export CLASS_DIRpath_to_class_images export OUTPUT_DIRpath_to_saved_model accelerate launch train_dreambooth.py \ --pretrained_model_name_or_path$MODEL_NAME \ --instance_data_dir$INSTANCE_DIR \ --class_data_dir$CLASS_DIR \ --output_dir$OUTPUT_DIR \ --with_prior_preservation --prior_loss_weight1.0 \ --instance_prompta photo of sks dog \ --class_prompta photo of dog \ --resolution512 \ --train_batch_size1 \ --sample_batch_size1 \ --gradient_accumulation_steps1 --gradient_checkpointing \ --learning_rate5e-6 \ --lr_schedulerconstant \ --lr_warmup_steps0 \ --num_class_images200 \ --max_train_steps800 \ --mixed_precisionfp169. 推理生成个性化图像训练完成后通过 [StableDiffusionPipeline] 并指定模型保存路径即可进行推理。请确保提示词中包含训练时使用的特殊标识符上例中的sksfrom diffusers import StableDiffusionPipeline import torch model_id path_to_saved_model pipe StableDiffusionPipeline.from_pretrained(model_id, dtypetorch.float16).to(cuda) prompt A photo of sks dog in a bucket image pipe(prompt, num_inference_steps50, guidance_scale7.5).images[0] image.save(dog-bucket.png)推理参数说明num_inference_steps50采样步数步数越多质量通常越高但耗时更长guidance_scale7.5无分类器引导强度CFG scale控制提示词对生成结果的约束程度。也可以按照第 7.3 节的方式从保存的中间训练检查点执行推理以便筛选出效果最好的中间产物。10. 补充脚本测试与运行验证仓库通过 examples/dreambooth/test_dreambooth.py 对训练脚本进行持续集成验证。测试使用hf-internal-testing/tiny-stable-diffusion-torch微型模型在极低分辨率--resolution 64和极短步数--max_train_steps 2下冒烟运行并断言输出目录中生成unet/diffusion_pytorch_model.safetensors与scheduler/scheduler_config.json等标准文件。这为读者提供了一种快速验证训练脚本可用性的方式——只需把模型换为 tiny 测试模型即可在数分钟内跑通全流程。结语DreamBooth 是个性化扩散模型微调的重要技术而 Diffusers 仓库中的 examples/dreambooth/train_dreambooth.py 提供了完整、可扩展的实现。通过本文你可以掌握基础 UNet 微调与先验保留损失的使用文本编码器联合微调需 ≥24GB VRAM与 LoRA 轻量微调检查点保存、恢复与推理的完整流程面向 16GB / 12GB / 8GB 显存的分级优化方案xFormers、梯度检查点、8 位 Adam、DeepSpeed。结合文中的参数表格与源码路径你可以按自己的数据集与硬件条件快速搭建一套可复现的 DreamBooth 个性化训练流水线。【免费下载链接】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),仅供参考
返回列表