当前位置:AIGC资讯 > AIGC > 正文

Stable Diffusion的微调方法原理总结

目录

1、Textural Inversion(简易)

2、DreamBooth(完整)

3、LoRA(灵巧)

4、ControlNet(彻底)

5、其他

1、Textural Inversion(简易)

        不改变网络结构,仅改变CLIP中token embedding的字典。在字典中新增一个伪词的embedding,fine-tune这个embedding的值。其他所有可调参数都冻结。

优点:训练量极小,需要的素材就是一张图。完全不改变神经网络中的任何参数。

缺点:效果一般。

TI的简洁激发了很多研究者的灵感,基于TI思路的研究出现了很多。

2、DreamBooth(完整)

        具体做法是,加入一个新词(sks)代表subject,embedding初始值继承原类型的词的embedding。调整了模型中全部可调参数,彻底的让模型学会subject。损失函数加入了监督功能,去监控漂移现象,防止灾难性遗忘“学会新的忘了旧的”。

在LoRA出现前,训练DreamBooth是潮流,但代价较大。

3、LoRA(灵巧)

        LoRA的网络是一种additional network,LoRA训练不改变基础模型的任何参数,只对附加网络内部参数进行调整。在生成图像时,附加网络输出与原网络输出融合,从而改变生成效果。

        由于LoRA是将矩阵压缩到低秩后训练,所以LoRA网络的参数量很小(千分之一),训练速度快。实验发现,低维矩阵对高维矩阵的替代损失不大。所以即便训练的矩阵小,训练效果仍然很好,已成为一种customization image generation范式。LoRA后来在结构上改进出不同的版本,例如LoHA,LyCORIS等。

LoRA详解:https://zhuanlan.zhihu.com/p/632159261

Self-Attention的LoRA微调代码:GitHub - owenliang/pytorch-diffusion: pytorch复现stable diffusion

代码分析:

用于替换的线性层 (Wq, Wk, Wv矩阵):

class CrossAttention(nn.Module):
    def __init__(self,channel,qsize,vsize,fsize,cls_emb_size):
        super().__init__()
        # Wq, Wk, Wv 矩阵使用LoRA微调降低参数量, W + WA * WB
        self.w_q=nn.Linear(channel,qsize)
        self.w_k=nn.Linear(cls_emb_size,qsize)
        self.w_v=nn.Linear(cls_emb_size,vsize)
        self.softmax=nn.Softmax(dim=-1)
        self.z_linear=nn.Linear(vsize,channel)
        self.norm1=nn.LayerNorm(channel)
        # feed-forward结构
        self.feedforward=nn.Sequential(
            nn.Linear(channel,fsize),
            nn.ReLU(),
            nn.Linear(fsize,channel)
        )
        self.norm2=nn.LayerNorm(channel)

找到模型中所有的Wq, Wk, Wv线性层并将其替换为Lora:

if __name__=='__main__':   # 加入LoRA微调的训练过程
    # 预训练模型
    model=torch.load('model.pt')

    # 向nn.Linear层注入Lora
    for name,layer in model.named_modules():
        name_cols=name.split('.')
        # 过滤出cross attention使用的linear权重
        filter_names=['w_q','w_k','w_v']
        if any(n in name_cols for n in filter_names) and isinstance(layer,nn.Linear):   # module名字中存在w_q, w_k, w_v且属于线性层
            # print(name)   # enc_convs.0.crossattn.w_q,enc_convs.0.crossattn.w_k,enc_convs.0.crossattn.w_v,……
            inject_lora(model,name,layer)

Lora具体实现与替换过程:

# Lora实现,封装linear,替换到父module里
class LoraLayer(nn.Module):
    def __init__(self,raw_linear,in_features,out_features,r,alpha):
        super().__init__()
        self.r=r   # 秩数
        self.alpha=alpha   # LoRA分支的权重比例系数
        self.lora_a=nn.Parameter(torch.empty((in_features,r)))   # 可训练参数
        self.lora_b=nn.Parameter(torch.zeros((r,out_features)))
    
        nn.init.kaiming_uniform_(self.lora_a,a=math.sqrt(5))   # WA 矩阵参数需要进行初始化

        self.raw_linear=raw_linear   # 原始模型权重 W
    
    def forward(self,x):    # x:(batch_size,in_features)
        raw_output=self.raw_linear(x)   
        lora_output=x@((self.lora_a@self.lora_b)*self.alpha/self.r)    # LoRA分支:x * (WA * WB * α/r)
        return raw_output+lora_output   # W + LoRA

def inject_lora(model,name,layer):
    name_cols=name.split('.')   # [enc_convs, 0, crossattn, w_q]

    # 逐层下探到linear归属的module
    children=name_cols[:-1]   # [enc_convs, 0, crossattn]
    cur_layer=model 
    for child in children:
        cur_layer=getattr(cur_layer,child)   # 逐层深入得到w_q, w_k, w_v层的属性
    
    #print(layer==getattr(cur_layer,name_cols[-1]))
    lora_layer=LoraLayer(layer,layer.in_features,layer.out_features,LORA_R,LORA_ALPHA)
    setattr(cur_layer,name_cols[-1],lora_layer)   # 把 crossattn 的 w_q/w_k/w_v层 的属性替换为LoraLayer

模型训练过程:冻结非Lora分支的所有参数

    # lora权重的加载
    try:
        restore_lora_state=torch.load('lora.pt')   # 加载训练好的Lora权重(lora_a, lora_b矩阵),enc_convs.0.crossattn.w_q.lora_a等
        model.load_state_dict(restore_lora_state,strict=False)
    except:
        pass 

    model=model.to(DEVICE)

    # 冻结非Lora参数
    for name,param in model.named_parameters():
        if name.split('.')[-1] not in ['lora_a','lora_b']:  # 非LoRA部分不计算梯度
            param.requires_grad=False
        else:
            param.requires_grad=True

模型推理过程:将Lora分支参数合并到原始模型参数中(相加)

if __name__=='__main__':
    # 加载模型
    model=torch.load('model.pt')

    USE_LORA=True

    if USE_LORA:   # 使用LoRA推理
        # 把Linear层替换为Lora
        for name,layer in model.named_modules():
            name_cols=name.split('.')
            # 过滤出cross attention使用的linear权重
            filter_names=['w_q','w_k','w_v']
            if any(n in name_cols for n in filter_names) and isinstance(layer,nn.Linear):
                inject_lora(model,name,layer)

        # lora权重的加载
        try:
            restore_lora_state=torch.load('lora.pt')
            model.load_state_dict(restore_lora_state,strict=False)
        except:
            pass 

        model=model.to(DEVICE)

        # lora权重合并到主模型(把LoRA权重加到原始模型权重中)
        for name,layer in model.named_modules():
            name_cols=name.split('.')

            if isinstance(layer,LoraLayer):   # 找到模型中所有的 LoraLayer 层
                children=name_cols[:-1]
                cur_layer=model 
                for child in children:
                    cur_layer=getattr(cur_layer,child)    # cur_layer = cross attention对象(包含修改过的wq, wk, wv)
                lora_weight=(layer.lora_a@layer.lora_b)*layer.alpha/layer.r   # 计算得到lora分支权重
                before_weight=layer.raw_linear.weight.clone()   # 原始模型权重W
                layer.raw_linear.weight=nn.Parameter(layer.raw_linear.weight.add(lora_weight.T)).to(DEVICE)    # 把Lora参数加到base model的linear weight上
                setattr(cur_layer,name_cols[-1],layer.raw_linear)   # 使用新的合并分支替换原来的两分支Lora结构

4、ControlNet(彻底)

        将神经网络快的不同权重,分别复制到“锁定”副本(locked copy)和“可训练”副本(trainable copy)中。按制定规则集成原图特征并生成新的内容,不会导致生成图和原图看起来毫无关系。

5、其他

Custom Diffusion基本建立在DreamBooth的基础上,通过消融实验证明了即使只训练交叉注意力层中的部分矩阵,也有非常好的fine-tune效果,不需要像DreamBooth那样全部参数调整。这种思路也引领了后续的一系列研究,但DreamBooth仍然是当时的范式。 与ControlNet同期有一种方法叫做T-2-l adapter,微调的参数更少,效果较CN差些,比CN发布晚了一点,被ControlNet的光芒遮挡了。 LORA的典型修改方案是LyCORIS,这个以二次元人物命名的方法把LoRA的思想应用在卷积层做改进,并且结合了一些其他算法进行了参数调整。 微调方法只是打包起来的tricks。模型建模研究是建构的过程,而不是发现的过程,有很大的自由度,不要被已有做法的说法限制自己的想象。

总结

这篇文章深入介绍了文本驱动生成领域中几种重要的图像微调(fine-tuning)与个性化方法,这些方法为用户或研究人员提供了灵活多变的定制方案,从而在大型基础模型之上生成更符合个人需求的内容。文章主要包含以下几大方面的内容:
### 1. **Textural Inversion(TI,简易版)**
- **方法简述**:仅通过调整CLIP模型中的token embedding字典来实现定制化,不需要改动神经网络的任何结构。通过在字典中增加一个伪词的embedding并进行微调来实现个性化生成。
- **优缺点**:优点在于训练数据量小、成本低、方法简便;但效果较为有限,不能达到高要求的个性化精度。
- **后续研究**:基于TI思路激发了大量研究者探索简易而高效个性化生成方式的灵感。
### 2. **DreamBooth(完整版)**
- **方法详述**:通过增加新的特定主题词汇(如sks)代表个性化对象,调整整个模型中可调参数的值,并用带监督的损失函数监控模型的学习效果,避免模型遗忘先前学习的通用内容。
- **应用影响**:在LoRA出现之前,DreamBooth是定制化生成的流行方案,尽管需要较大的计算资源。
### 3. **LoRA(灵巧版)**
- **核心理念**:作为additional network引入,不直接改变基模型的参数,通过调整少量低秩矩阵的参数影响生成结果。由于训练矩阵低维,所需资源极少但效果佳,快速成为主流的customization图像生成方式。
- **技术创新**:后续出现多个结构改进的版本如LoHA、LyCORIS,体现了其广泛的应用性和发展潜力。
### 4. **ControlNet(彻底版)**
- **工作原理**:通过复制基础神经网络层的权重至两个版本——锁定的不可变版和可训练的调优版,控制网络如何融合原始图片信息和指令文本信息以生成个性化图像,避免生成的图片与原始素材相脱节。
### 5. **其他研究与实践**
- **Custom Diffusion与T-2-I adapter**:前者在DreamBooth基础上进行更精简的调整尝试;后者是一种轻量级参数调整的方案,尽管性能不如ControlNet但仍为研究提供了新的方向。
- **LoRA扩展——LyCORIS**:该方法结合了LoRA思路在卷积层的改进和更多元的参数优化策略,提升了模型的适用范围和精度。
### 结语
这些微调方法展现了生成模型在应用层面丰富的多样性和灵活的定制性,虽然基于各种技巧的组合与优化,但它们都为研究和实际应用提供了重要的基础和创新路径。这些方案的发展,尤其是LoRA与ControlNet,将进一步推动AI艺术创作与内容个性化的发展进程。

更新时间 2024-08-31