Stable Diffusion 微调及推理优化
腾
云煮鸡约 8 分钟
云煮鸡约 8 分钟本文作者：刘远，腾讯云泛互联网首席解决方案架构师。
概述
Stable Diffusion 微调
Stable Diffusion 微调的目标，是将新概念注入预训练模型，利用新注入的概念以及模型的先验知识，基于文本引导条件生成自定义图片。目前主流训练 Stable Diffusion 模型的方法有 Full FineTune、Dreambooth、Text Inversion 和 LoRA，不同方法的实现逻辑和使用场景不同，选型简单对比如下：
| 训练方法 | 方法 | 局限性 |
|---|---|---|
Text Inversion | 使用提供的一组图片训练一个新单词的Embedding ，并将其与词汇表中的已有单词关联起来，这个新单词即为这组图片概念的指代。 | 训练过程只对应 Embedding，扩散模型没有新知识输入，所以也无法产生新的内容。 |
Full FineTune | 最朴素的方式，使用图片+ 标注的数据集，进行迭代训练，数据集标注可以选择BLIP来生成。训练直接对原模型的所有权重进行调整。 | 容易过拟合，导致生成图片的多样性不够，结果难以控制。模型体积大，不便于传播。 |
Dreambooth | 提供代表某个新概念（instance） 对应的一组图像，并使用罕见字符（identifier） 进行概念Mapping，训练过程充分考虑原有相关主题（class）生成，避免过拟合。训练直接对原模型的所有权重进行调整。 | 训练过程只针对新概念 （instance），多样性差。如果需要多概念生成，需要多次训练。模型体积大，不便于传播。 |
LoRA（w Dreambooth） | 冻结预训练模型参数，在每个Transformer块插入可训练层，不需要完整调整 UNet 模型的全部参数。训练结果只保留新增的网络层，模型体积小。 | 训练效果不如Dreambooth |
需要注意的是，LoRA 是一种加速训练的方法，Stable Diffusion 从大语言模型微调中借鉴而来，可以搭配 Full FineTune 或 Dreambooth 使用。针对上述几种训练方法，我们在 A10-24G 机型上进行测试，5-10张训练图片，所需资源和时长对比如下：
| Dreambooth | Text Inversion | LoRA | |
|---|---|---|---|
模型体积 | ~ 4 GB | 3-10 KB | 3-10 MB |
| 训练显存 | ~ 22GB | ~ 5GB | ~ 20GB |
| 训练时长 | 15分钟 | 1小时 | 10分钟（w Dreambooth） 4小时（w Full FineTune） |
Dreambooth
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_weight=1.0 \
--mixed_precision=fp16 \
--instance_prompt="a photo of az baby" \
--class_prompt="a photo of baby" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--num_class_images=200 \
--max_train_steps=800
python ../scripts/convert_diffusers_to_original_stable_diffusion.py --model_path ./dreambooth_baby --checkpoint_path dreambooth_baby.safetensors --use_safetensors
LoRA（w Dreambooth）
accelerate launch train_dreambooth_lora.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--class_data_dir=$CLASS_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a photo of az baby" \
--class_prompt="a photo of baby" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--checkpointing_steps=100 \
--learning_rate=1e-4 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=1000 \
--validation_prompt="a photo of az baby" \
--validation_epochs=50 \
--seed="0"
python diffusers-lora-to-safetensors.py --file pytorch_lora_weights.bin
Stable Diffusion 性能优化
Dreambooth 优化
docker run -it --gpus=all --network=host -v /[diffusers_model_directory]:/[custom_container_directory] sd_taco:v3 bash
script_model = torch.jit.trace(model, test_data, strict=False)
script_model.save("trace_module.pt")
import torch
import taco
import ostaco_path = os.path.dirname(taco.__file__)
torch.ops.load_library(os.path.join(taco_path, "torch_tensorrt/lib/libtorchtrt.so"))
optimized_model = torch.jit.load("optimized_recursive_script_module.pt")
pic = torch.rand(1, 4, 64, 64).cuda() // picture
timesteps = torch.tensor([1]*1) // timesteps
context = torch.randn(1, 77, 768) // text embedding
with torch.no_grad():
output = optimized_model(pic, timesteps, context)
print(output)
LoRA 优化
python networks/merge_lora.py --sd_model ../v1-5-pruned-emaonly.safetensors --save_to ../lora-v1-5-pruned-emaonly.safetensors --models <LoRA文件目录> --ratios <LoRA权重>