Diffusers代码学习-LoRA训练

LoRA（Low-Rank Adaptation of Large Language Models）是一种流行的轻量级训练技术，它显著减少了可训练参数的数量。它的工作原理是在模型中插入少量的新权重，并且只训练这些权重。这使得使用LoRA进行训练的速度更快、内存高效，并产生更小的模型权重（几百MB），更容易存储和共享。LoRA还可以与DreamBooth等其他训练技术相结合，以加快训练速度。

1. 下载源代码及安装。

在运行脚本之前，请确保从以下源安装库：

git clone https://github.com/huggingface/diffusers

cd diffusers

pip install .

导航到包含train_dreambooth.py脚本的示例文件夹，并为正在使用的脚本安装所需的依赖项：

cd examples/dreambooth

pip install -r requirements.txt

注意需要修改train_dreambooth.py，将accelerator_project_config改为project_config。否则会报错returned non-zero exit status 1.

2. 安装及注册wandb

pip install wandb

登录Weights & Biases: The AI Developer Platform (wandb.ai)

3. 准备训练数据集

同时下载好训练所用的数据集。

import os

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

from huggingface_hub import snapshot_download

local_dir = "./dog"

snapshot_download(

    "diffusers/dog-example",

    local_dir=local_dir,

    repo_type="dataset",

    ignore_patterns=".gitattributes",

)

注意如果下载的文件中除了图片以外还有其他目录比如.huggingface目录，需要删掉。否则会报错PermissionError: [Errno 13] Permission denied

4. 初始化accelerate

然后需要用到accelerate库，可帮助在多个GPU/TPU上或以混合精度进行训练。它将根据硬件和环境自动配置训练设置。

这里不选择任何配置，使用设置默认值运行accelerate库

accelerate config default

5. 运行训练脚本

accelerate launch train_dreambooth_lora.py 

--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" 

--instance_data_dir="dog" 

--output_dir="model_dog" 

--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" 

6. 查看生成结果

7. 登录wandb查看训练过程