用LoRA训练Stable Diffusion时,首先冻结模型的权重,然后在U-Net结构中注入LoRA模型,将其与交叉注意力模块结合在一起,最后微调时将只对这部分参数进行更新。
数据收集
本次微调将使用数码宝贝数据集作为下游细分任务,数据来源于数码兽数据库(http://digimons.net/digimon/chn.html)。Stable Diffusion的训练数据非常直观,就是一张图片对应一段文本描述,因此我们需要先通过一个爬虫将数据整理下来。
我们需要从数码兽图鉴网页中爬取所有数码兽的信息,包括名称、介绍和对应的图片链接。然后将这些信息整理成下面这种格式,保存到对应的文件夹中。
# 数据格式 metadata.jsonl + 图片
folder/train/metadata.jsonl #存储caption描述
folder/train/0001.png
folder/train/0002.png
folder/train/0003.png
# metadata.jsonl中的内容
{"file_name": "0001.png", "text": "image 1 description"}
{"file_name": "0002.png", "text": "image 2 description"}
{"file_name": "0003.png", "text": "image 3 description"}
为了实现这个任务,我们需要使用Python的爬虫和文件操作相关的库。首先,使用requests库获取数码兽图鉴页面的HTML内容,并使用BeautifulSoup库解析HTML,以便对页面进行提取信息。然后我们分析这个页面,所有的数码兽都存在一个id为digimon_list的ul列表中,每一行就是一个li标签,这个标签里面有一个a标签,对应了该数码兽的详情链接。
接下来,我们遍历页面中所有的li标签,提取数码兽的名称和详情页面链接。然后进入详情页面,获取数码兽的介绍和图片链接。最后,将这些信息整理成指定的格式,并保存到对应的文件夹中。具体而言,我们需要在指定的文件夹中创建一个metadata.jsonl文件来保存每个图片的文件名和对应的描述文本,并使用文件名对应的顺序来保存对应的图片文件。
最终,我们会得到一个数据集,其中包含每个数码兽的名称、介绍和对应的图片,以及一个metadata.jsonl文件,其中保存了每个图片的文件名和对应的描述文本。
python
import os
import json
import requests
from bs4 import BeautifulSoup
# 创建文件夹
data_dir = "./train"
if not os.path.exists(data_dir):
os.makedirs(data_dir)
# 请求数码兽图鉴页面
url = "http://digimons.net/digimon/chn.html"
response = requests.get(url)
soup = BeautifulSoup(response.content, "html.parser")
# 遍历所有的 li 标签
digimon_list = soup.find("ul", id="digimon_list")
for digimon in digimon_list.find_all("li"):
try:
# 获取数码兽名称和详情页面链接
name = digimon.find('a')["href"].split('/')[0]
detail_url = "http://digimons.net/digimon/" + digimon.find('a')["href"]
print(f"detail_url: {detail_url}")
# 进入详情页面,获取数码兽介绍和图片链接
response = requests.get(detail_url)
soup = BeautifulSoup(response.content, "html.parser")
caption = soup.find("div", class_="profile_eng").find('p').text.strip()
img_url = f"http://digimons.net/digimon/{name}/{name}.jpg"
# 保存图片
img_data = requests.get(img_url).content
file_name = f"{len(os.listdir(data_dir)) + 1:04d}.png"
with open(os.path.join(data_dir, file_name), "wb") as f:
f.write(img_data)
# 将数据整理成指定的格式,并保存到对应的文件中
metadata = {"file_name": file_name, "text": f"{name}. {caption}"}
with open(os.path.join(data_dir, "metadata.jsonl"), 'a') as f:
f.write(json.dumps(metadata) + '\n')
except Exception as _:
pass
数据集整理完成后,如果我们想让其他人也可以用我们的数据集,那么就可以将其发布到Hugging Face Hub上。Hugging Face Hub是一个收集了多个领域中多种任务的模型以及数据集的平台,使用户可以非常方便地下载和使用各种资源。并且,Hugging Face也非常鼓励用户上传自己的数据集,以帮助壮大AI学习社区并加快其发展步伐,造福大家。
如果你还没有Hugging Face Hub账号,需要先去其官网注册一个账号,然后按照以下步骤创建数据集。
- 点击页面右上角profile下的New DataSet选项创建一个新的数据集仓库,如图7-4所示。
- 输入数据集的名称,并选择是否为公开数据集,公开数据集对所有人可见,而私有数据集仅对自己或组织成员可见。
图7-4 Hugging Face创建数据集页面样例
1)进入数据集页面,点击顶部菜单栏的“Files and versions”,然后点击右边“Add file”按钮下的“upload files”按钮,之后将图片文件和训练数据元文件直接拖拽到上传文件框,最后编写修改信息,点击提交即可,如图7-5所示。
图7-5 Hugging Face上传数据集内容页面样例
训练参数设置
需要注意的是,为了保证能够成功运行最新版的训练代码,建议通过源码重新安装Diffusers库。
bash
pip install git+https://github.com/Hugging Face/diffusers
然后我们需要初始化Accelerate分布式训练环境。
bash
> accelerate config
[2023-07-20 18:37:53,537] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
NOTE: Redirects are currently not supported in Windows or MacOs.
--------------------------------------------------------------------------------------------------------------------In which compute environment are you running?
This machine
--------------------------------------------------------------------------------------------------------------------Which type of machine are you using?
No distributed training
Do you want to run your training on CPU only (even if a GPU / Apple Silicon device is available)? [yes/NO]: NO
Do you wish to optimize your script with torch dynamo?[yes/NO]: NO
Do you want to use DeepSpeed? [yes/NO]: NO
What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]: all
--------------------------------------------------------------------------------------------------------------------Do you wish to use FP16 or BF16 (mixed precision)?
no
accelerate configuration saved at C:\Users\admin/.cache\Hugging Face\accelerate\default_config.yaml
模型训练与测试
当前,LoRA技术主要支持 UNet2DConditionalModel。Diffusers团队已经推出了一款适用于LoRA的微调脚本,这款脚本的优势在于它能够在仅有11GB GPU RAM的环境中稳定运行,而且不需要依赖8-bit等优化技术。下面我们将展示了如何使用此脚本结合数码兽数据集来进行模型的微调操作。
bash
accelerate launch --mixed_precisinotallow="fp16" train_text_to_image_lora.py \
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
--train_data_dir="./train_data" \
--dataloader_num_workers=0 \
--resolutinotallow=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="./finetune/lora/digimon" \
--checkpointing_steps=500 \
--validation_prompt="Blue Agumon" \
--seed=1024
该脚本是启动了一个混合精度为fp16的加速微调训练任务。它采用的预训练模型是runwayml/stable-diffusion-v1-5,并从"./train_data"路径下获取训练数据。脚本配置为单线程加载数据,并将图像解析为512x512的分辨率,同时允许中心裁剪和随机翻转。虽然每次只处理一个批次的数据,但它会累计四个批次的梯度进行一次更新。训练的最大步数被设置为15000步,学习率为0.0001,并限制了梯度的最大范数为1。学习率调度器采用的是余弦退
火策略,不进行预热。训练结果将保存在"./finetune/lora/digimon"目录下,每500步保存一次检查点。此外,验证的提示词设置为"Blue Agumon",并指定了随机种子为1024,以确保实验的可重复性。
正如我们前面所讨论的,LoRA的主要优势之一是可以通过训练比原始模型小几个数量级的权重来获得出色的结果。通过load_attn_procs函数,我们可以在原始的Stable Diffusion模型权重之上加载额外的权重。
python
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
model_path = "runwayml/stable-diffusion-v1-5"
LoRA_path = "./finetune/lora/digimon" # 修改成本地LoRA模型路径
pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.unet.load_attn_procs(LoRA_path)
pipe.to("cuda")
image = pipe("blue skin agumon", num_inference_steps=50).images[0]
image.save("test.png")
在上面这段代码中,我们首先定义了两个路径,一个是主模型路径model_path,如果还是使用原生的Stable Diffusion则不需要修改,还有一个是LoRA模型的路径LoRA_path,需要将LoRA_path修改为正确的本地LoRA模型路径。接下来创建StableDiffusionPipeline流水线对象,然后通过load_attn_procs方法用于加载本地的LoRA模型,并将其应用于流水线的unet模块,再将管道移至GPU以加快推理速度。最后,我们给定一个提示词“blue skin agumon”,让模型生成一个蓝色皮肤的亚古兽,训练数据集的亚古兽图片与生成的图片如图7-6所示。
图7-6 原始亚古兽图片(左)与微调后的模型生成的“蓝色亚古兽”图片(右)对比
在推理时我们还可以调整LoRA的权重系统:
图片
如果将设置为0时,其效果与只使用主模型完全一致;如果将设置为1时,与使用的效果相同。因此,如果LoRA存在过拟合的现象,我们可以将设置为较低的值,如果使用LoRA的效果不太明显,那我们可以将设置为略高于1的值。
除了使用单个LoRA模型,还可以将多个LoRA模型同时添加到一个主模型中,同样也可以调整两个LoRA模型的权重:
图片
如果将和都设置为0.5,那么在主模型添加的就是两个LoRA模型的平均权重。如果将设置为0.7,设置为0.3,那么第一个LoRA模型的效果将占据70%的效果。
在代码中,将LoRA权重与冻结的预训练模型权重合并时,也可以选择调整与参数合并的权重数量scale,scale值为0表示不使用LoRA权重,而scale值为1表示使用完全微调的 LoRA 权重。
python
pipe.unet.load_attn_procs(lora_model_path)
pipe.to("cuda")
image = pipe(
"A agumon with blue skin.", num_inference_steps=25, guidance_scale=7.5, cross_attention_kwargs={"scale": 0.5}
).images[0]
image.save("blue_pokemon.png")