2023-08-20 14:01 算法工程师

关注

深入浅出解析LoRA完整核心基础知识 | 【算法兵器谱】

Rocky Ding

公众号：WeThinkIn

写在前面

【算法兵器谱】栏目专注分享AI行业中的前沿/经典/必备的模型&论文，并对具备划时代意义的模型&论文进行全方位系统的解析，比如Rocky之前出品的爆款文章Make YOLO Great Again系列。也欢迎大家提出宝贵的优化建议，一起交流学习💪

大家好，我是Rocky。

在AIGC时代到来之后，LoRA无疑成为了AI绘画领域与Stable Diffusion（SD）配合最多的模型，SD模型+LoRA的组合，生成了一个又一个精美的图片与视频，创造了很多脑洞大开的风格，大大提高了AI生成内容的质量与效率。

Rocky在本文将会把LoRA最核心的知识分享给大家，包括LoRA的基本原理、流程训练、使用方法以及一些“经验”感悟，力求能让大家零基础入门使用LoRA进行AI内容创造！

在这里也预告一下，下一期Rocky将分享LoRA的进阶用法，更多干货思考分享，欢迎大家一键三连！

关于作者

Rocky在校招期间拿到了北上广深杭等地的约10个算法offer，现在是一名算法研究员，目前专注于AIGC创新产品的落地应用以及AI算法解决方案的商用。

在研究生期间，Rocky曾在京东研究院，星环科技，联想研究院，北大方正信产集团研究院，百融云创等公司做算法实习生，对不同性质公司的商业闭环逻辑比较了解。

Rocky喜欢参加算法竞赛，多次获得CVPR，AAAI，Kaggle等顶级平台的算法竞赛冠军和Top成绩。

Rocky相信人工智能，数据科学，商业逻辑，金融工具，终身成长，以及顺应时代的潮流会赋予我们超能力。

Rocky是AI人才发展联盟（AIHIA）的联合创始人，积极在业余时间进行创业实践与基本面扩展。

Rocky喜欢分享和交流，秉持着“也要学习也要酷”的生活态度，希望能和大家多多交流。AI算法，面试，简历，求职等问题都可直接和我交流～

So，enjoy：

正文开始

----【目录先行】----

LoRA的核心基础知识
LoRA的优势
轻松入门LoRA训练
使用LoRA轻松生成图片
LoRA的训练技巧

【一】LoRA的核心基础知识

在AIGC爆发后，大模型进入了业界的视野，大模型强大的能力让AI彻底破圈，可谓是天下谁人不识Stable Diffusion和ChatGPT。

但是大模型参数量巨大，训练成本较高，当遇到一些下游细分任务时，对大模型进行全参训练性价比不高，同时这些下游细分任务的域比较好约束，在这种情况下，本文的主角——LoRA就出场了。

我们可以使用SD模型+LoRA微调训练的方式，只训练参数量很小的LoRA，就能在下游细分任务中取得不错的效果。

LoRA的训练逻辑是首先冻结SD模型的权重，然后在SD模型的U-Net结构中注入LoRA模块，将其并与crossattention模块结合，并只对这部分参数进行微调训练。

也就是说，对于SD模型权重 $W_{o} \in \mathbb{R}^{n\times m}$ ，我们不再对其进行全参微调训练，我们对权重加入残差的形式，通过训练 $\Delta w$ 来完成优化过程：

$W' = W \Delta W$

其中 $\Delta w = AB^{T},A \in \mathbb{R}^{n\times d} ,B \in \mathbb{N}^{d\times m}, d \ll n$ ，其是由两个低秩矩阵的乘积组成。由于下游细分任务的域非常小，所以 $d$ 可以取得很小，很多时候我们可以取 $1$ 。因此在训练完成之后，我们可以获得一个参数量远小于SD模型的LoRA模型。

通常来说，对于矩阵 $A$ ，我们使用随机高斯分布初始化，并对于矩阵 $B$ 使用全 $0$ 初始化，使得在初始状态下这两个矩阵相乘的结果为 $0$ 。这样能够保证在初始阶段时，只有SD模型（主模型）生效。

LoRA主要与主模型的Crossattention结合训练

经过开源生态的持续验证，我们发现使用LoRA进行微调与SD模型全参微调的效果相差不大。

Rocky认为LoRA成为了AIGC时代的“残差模块”，SD模型的“得力助手”。

LoRA大幅降低了SD模型训练时的显存占用，因为并不优化主模型（SD模型），所以主模型对应的优化器参数不需要存储。但计算量没有明显变化，因为LoRA是在主模型的全参梯度基础上增加了“残差”梯度，同时节省了主模型优化器更新权重的过程。

与此同时，在推理过程中，由于可以直接把LoRA合并到主模型中，从而整个推理成本是没有增加的，这无疑让LoRA有了更多的实用价值。

【二】LoRA的优势

与SD模型全参训练相比，LoRA训练速度更快。
非常低的算力要求。我们可以在2080Ti级别的算力设备上进行LoRA训练。
由于只是与SD模型的结合训练，LoRA本身的参数量非常小，最小可至3M左右。
能在小数据集上进行训练（10-30张），并且与不同的SD模型都能较好的兼容与迁移适配。
训练时主模型参数保持不变，LoRA能更好的在主模型的能力上优化学习。

【三】轻松入门LoRA训练

我们可以直接使用diffusers库里的代码进行LoRA训练，关于diffusers与Stable Diffusion，大家可以阅读Rocky之前的文章：深入浅出解析Stable Diffusion完整核心基础知识 |【算法兵器谱】

本次Rocky选用个人非常喜欢的宝可梦数据集作为下游细分任务，并让LoRA能够学习到所有宝可梦的特征，并能生成我们想要的独一无二的“新宝可梦”！

我们首先需要准备数据，包括图片和对应的caption：

数据集格式：图片+caption

由于本次的宝可梦数据集已经封装好了，所以我们可以直接调用数据集路径即可，若是我们想要训练自己的数据集，需要制作diffusers能够读取的数据格式：

#数据格式 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": "This is a golden retriever playing with a ball"}
{"file_name": "0002.png", "text": "A german shepherd"}
{"file_name": "0003.png", "text": "One chihuahua"}

好的，完成了数据集的制作，我们就可以进行训练了，我们可以使用diffusers库中的train_text_to_image_lora.py作为训练代码，并传入相关的预设参数：

export MODEL_NAME="runwayml/stable-diffusion-v1-5" #选用的主模型
export OUTPUT_DIR="/lora/pokemon" #LoRA模型保存地址
export DATASET_NAME="lambdalabs/pokemon-blip-captions" #训练数据路径

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 \ #Batch-Size
  --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 \
  --checkpointing_steps=500 \
  --seed=2048

到这里，我们就完成了LoRA的训练流程。

【四】使用LoRA轻松生成图片

完成LoRA的训练后，我们可以使用LoRA进行推理，生成我们想要的图片：

#读取diffuers库
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

#设置SD模型路径和LoRA模型路径
model_path = "runwayml/stable-diffusion-v1-5" #修改成本地主模型路径
LoRA_Path = "WeThinkIn" #修改成本地LoRA模型路径

#初始化SD模型，加载预训练权重
pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

# 加载LoRA weights ～3 MB
pipe.unet.load_attn_procs(LoRA_Path)
pipe.to("cuda")

#接下来，我们就可以运行pipeline了
image = pipe("blue pokemon", num_inference_steps=25).images[0]
image.save("test.png")

目前为止，我们已经使用LoRA完成常规的图像生成流程了。

不过除了主模型+LoRA的形式，我们还可以调整LoRA的权重：

$W' = W \alpha\Delta W$

下面的图示展示了将 $\alpha$ 从 $0$ 缩放到 $1$ 对图像生成产生的影响。

将 $\alpha$ 设置为 $0$ 时，与只使用主模型的效果完全相同；将 $\alpha$ 设置为 $1$ 时，与使用 $W' = W \Delta W$ 的效果相同。如果LoRA存在过拟合的情况，我们可以将 $\alpha$ 设置为较低的值。如果使用LoRA的效果不太明显，拿我们可以将 $\alpha$ 设置为略高于 $1$ 的值。改变权产生的效果