当前位置:   article > 正文

stable diffusion 1.x 模型训练概述_stable diffusion冻结训练

stable diffusion冻结训练

本文主要介绍sd发展的第一阶段版本的模型架构和一些微调训练方法,后续版本SDXL,LCM 版本再写文章继续介绍

SD 1.x 版本模型结构

autoencoder(VAE):encoder将图像压缩到latent空间,而decoder将latent解码为图像;

CLIP text encoder:提取输入text的text embeddings,通过cross attention方式送入扩散模型的UNet中作为condition;

UNet:扩散模型的主体,用来实现文本引导下的latent生成。

VAE

autoencoder是一个基于encoder-decoder架构的图像压缩模型,对于一个大小为H×W×3的输入图像,encoder模块将其编码为一个大小为h× w × c的latent,其中f = H/h = W/h为下采样

率(downsampling factor)。

使用diffusers库来加载autoencoder模型,并使用autoencoder来实现图像的压缩和重建,代码如下所示:

使用两张图片在256x256和512x512下的重建效果对比

如下所示,第一列为原始图片,第二列为512x512尺寸下的重建图,第三列为256x256尺寸下的重建图。对比可以看出,autoencoder将图片压缩到latent后再重建其实是有损的,比如会出现文字和人脸的畸变,在256x256分辨率下是比较明显的,512x512下效果会好很多。

CLIP

CLIP的英文全称是Contrastive Language-Image Pre-training,即一种基于对比文本-图像对的预训练方法或者模型。CLIP是一种基于对比学习的多模态模型

与CV中的一些对比学习方法如moco和simclr不同的是,CLIP的训练数据是文本-图像对:一张图像和它对应的文本描述,这里希望通过对比学习,模型能够学习到文本-图像对的匹配关系。

如下图所示,CLIP包括两个模型:Text Encoder和Image Encoder,其中Text Encoder用来提取文本的特征,可以采用NLP中常用的text transformer模型;而Image Encoder用来提取图像的特征,可以采用常用CNN模型或者vision transformer。

CLIP text encoder

SD采用CLIP text encoder来对输入text提取text embeddings,具体的是SD 1.X 采用目前OpenAI所开源的最大CLIP模型:clip-vit-large-patch14,这个CLIP的text encoder是一个transformer模型(只有encoder模块):层数为12,特征维度为768,模型参数大小是123M。

对于输入text,送入CLIP text encoder后得到最后的hidden states(即最后一个transformer block得到的特征),其特征维度大小为77x768(77是token的数量),这个细粒度的text embeddings将以cross attention的方式送入UNet中。在transofmers库中,可以如下使用CLIP text encoder:

Unet

SD的扩散模型是一个860M的UNet,其主要结构如下图所示(这里以输入的latent为64x64x4维度为例),其中encoder部分包括3个CrossAttnDownBlock2D模块和1个DownBlock2D模块,而decoder部分包括1个UpBlock2D模块和3个CrossAttnUpBlock2D模块,中间还有一个UNetMidBlock2DCrossAttn模块。encoder和decoder两个部分是完全对应的,中间存在skip connection。注意3个CrossAttnDownBlock2D模块最后均有一个2x的downsample操作,而DownBlock2D模块是不包含下采样的。

其中CrossAttnDownBlock2D模块的主要结构如下图所示,text condition将通过CrossAttention模块嵌入进来,此时Attention的query是UNet的中间特征,而key和value则是text embeddings。

SD和DDPM一样采用预测noise的方法来训练UNet,其训练损失也和DDPM一样,这里的C为text embeddings,此时的模型是一个条件扩散模型

CFG

在训练条件扩散模型时,往往会采用Classifier-Free Guidance(这里简称为CFG),所谓的CFG简单来说就是在训练条件扩散模型的同时也训练一个无条件的扩散模型,同时在采样阶段将条件控制下预测的噪音和无条件下的预测噪音组合在一起来确定最终的噪音,具体的计算公式如下所示:

  1. 这里的W为guidance scale,当W越大时,condition起的作用越大,即生成的图像其更和输入文本一致。 
  2. CFG的具体实现非常简单,在训练过程中,我们只需要以一定的概率(比如10%)随机drop掉text即可,这里我们可以将text置为空字符串(此时依然能够提取text embeddings)。

  3. CFG 中的t 参数就是我们经常使用negative_promt,可以让模型不生成负面提示的内容

  4. CFG对于提升条件扩散模型的图像生成效果是至关重要的

CFG背后的技术原理,感兴趣的可以阅读CFG的论文:Classifier-Free Diffusion Guidance以及guided diffusion的论文Diffusion Models Beat GANs on Image Synthesis。


模型微调

text to image training

sd最早的训练版本可以参考sd的核心论文 LDM的作者使用的开发训练ldm的仓库版本:

GitHub - pesser/stable-diffusion

也可以参考compvis论文的官方版本:GitHub - CompVis/stable-diffusion: A latent text-to-image diffusion model

diffusers 库对整个训练和推理都做了很好的封装,可以参考官方的代码仓库

https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img
下面以diffusers库的模型训练为例

diffusers/examples/text_to_image/train_text_to_image.py at main · huggingface/diffusers · GitHub

加载模型组件

模型训练

textual inversion (embeding)

  • 工作原理

  • 文本嵌入和反转过程概要

包含占位符字的字符串首先被转换为token(即字典中的单词或子单词索引)。这些token被转换为连续向量表示(“embeding”,v)。

最后,将嵌入向量转换成引导生成模型的单一条件代码co(y)。不断优化嵌入向量v,与我们的构造词S'相关联,使用目标函数构建新的loss。

基于Textual Inversion技术,仅使用3-5张用户提供的一组图像,在冻结了网络权重的sd模型的文本嵌入空间中使用新的”单词”去表征图像的信息,这些单词”能够组合到自然语言中,去让用户直观的控制图像的生成内容,从而达到个性化的创作

  • 特点

模型文件小,~30KB。通常不能捕获物品细节,擅长风格转换。可在Prompt中同时使用多个Embeddings

  • 核心训练代码

  • 案例

在新的场景中嵌入用户指定的概念风格

使用prompt生成物体的变体形状

gta5 训练图片

原始模型生成的图片

a sports car in street

加入训练的embeding模型生成的图片
a <gta5-artwork> style sports car in street

潜在的问题

出不来效果,欠拟合

效果不是很稳定

Dreambooth

背景

给定一个特定的物体,使用当前最好的文本到图像模型在不同的场景中生成它,同时保持其关键视觉特征的高保真度是非常具有挑战性的

新的方法

Dreambooth,可以通过学习一个主题的几张(3~5)图像对SD模型进行微调,来生成新的图像。

原理

给定 ∼ 3 − 5 个主题图像,我们微调文本到图像的扩散模型,包含两个步骤:

(a) 将输入图像与文本提示配对,微调低分辨率文本到图像模型,包含唯一标识符和主题词所属类别的名称(例如,“A [V] 狗”),同时,

应用了class-specific的先验损失,这样就可以使用之前类所属的语义信息,可以让模型更好的使用类名(“A dog”)生成属于该主题的不同实例

(b) 微调超分辨率组件与低分辨率和从我们的输入图像集中获取的高分辨率图像,这使我们能够保持高保真度到小主题的细节

用户输入一组图像包含一个主题,可以是人,车等任意物体,给预训练的模型注入新的自定义主题,在sd模型的文本编码器和Unet中学习新的主题进行微调模型,在模型中使用了唯一标识符与特定主题词进行了绑定,这样对包含主题的内容进行重新背景化、文本引导的视图合成、外观修改和艺术渲染(同时保留主题的关键特征)等操作

案例

文本引导的视图合成

外观修改

艺术渲染

装饰

Lora

概要

LoRA: Low-Rank Adaptation of Large Language Models 是微软研究员引入的一项新技术,主要用于处理大模型微调的问题。

目前超过数十亿以上参数的具有强能力的大模型 (例如 GPT-3) 通常在为了适应其下游任务的微调中会呈现出巨大开销。

LoRA 建议冻结预训练模型的权重并在每个 Transformer 块中注入可训练层 (秩-分解矩阵)。

因为不需要为大多数模型权重计算梯度,所以大大减少了需要训练参数的数量并且降低了 GPU 的内存要求。

这个方法的重点在于:预训练权重不需训练,因此没有梯度,仅训练低秩矩阵那部分的参数。

研究人员发现,通过聚焦大模型的 Transformer 注意力块,使用 LoRA 进行的微调质量与全模型微调相当,同时速度更快且需要更少的计算

lora方法论文链接:https://arxiv.org/abs/2106.09685

lora 实现代码:GitHub - microsoft/LoRA: Code for loralib, an implementation of "LoRA: Low-Rank Adaptation of Large Language Models"

原理

权重更新的部分

其中W0就是更新前的权重(在一开始就是预训练权重),而 △W 就是更新的那部分,也就是经过反向传播得到梯度后计算出来的需要更新的量。

对模型的哪些部分做低秩分解

可是,在如今的 202x 年代,模型通常有 N 个权重矩阵,那么应该对其中的哪些做低秩分解呢?还是说,应该暴力地、一个不拉地通杀?

对于这个问题,论文作者选择了简单的方式,他仅将 LoRA 应用于 self-attention 层中的 projection matrices,而其中的 MLP 模块以及 self-attention 层以外的结构均视而不见了。

sd lora

最早是社区一位大神实现的lora版本:

GitHub - cloneofsimo/lora: Using Low-rank adaptation to quickly fine-tune diffusion models.

这个lora方法后来也被集成到了diffusers 库中了

diffusers lora方法可以参考以下文档,如何训练和使用官方文档都有很详细的说明:

https://huggingface.co/docs/diffusers/training/lora#low-rank-adaptation-of-large-language-models-lorahttps://huggingface.co/blog/zh/lora

https://www.wpsshop.cn/w/从前慢现在也慢/article/detail/952106

推荐阅读
相关标签
  

闽ICP备14008679号