理解 Stable Diffusion UNet 网络

2024-5-26

前面的学习中,我们把 SD UNet 网络当成黑盒,不太影响对图片生成大致原理的理解,但在继续学 SD 的过程中,发现 ControlNet、T2I-Adapter、IPAdapter 等这些技术,都是在原 SD 网络模型上以各种方式对网络做修改叠加,要理解这些技术,还是得先了解下 SD UNet 网络结构的一些细节,不然看得很费劲。

SD 模型构成

从之前的学习我们知道,Stable Diffusion 模型里包含了三个组件:CLIP、VAE、UNet,这三个组件的参数和大小分布(来源):

组件 参数个数 文件大小 占比
CLIP 123,060,480 492 MB 12%
VAE 83,653,863 335 MB 8%
UNet 859,520,964 3.44 GB 80%
Total 1,066,235,307 4.27 GB 100%

整个生图的核心网络就是 UNet。UNet 最初是用于医学图像分割任务而提出来的,但它的特性展现了在图像其他领域的应用潜力,后续经过扩散模型的改进,很好应用在了图像生成上,所以 Stable Diffusion 的 UNet 实际上在原 UNet 网络架构上做了一些改造。

基础 UNet 网络

我们先来看看原 UNet 网络架构:

1

  1. 左边输入图片,经过整个网络处理,右边输出同尺寸图片。(原 UNet 网络用于医学图像识别分割,所以图上右边标的输出是一张同尺寸分割图。SD 这里的输出是降噪图)
  2. 左边下采样(也可以称为编码器),右边上采样(也可以称为解码器),一张图片经过一层层下采样计算,尺寸逐渐减小(图中的网络是减小到32×32),再经过右边层层上采样,恢复到原尺寸。那这里下采样和上采样的作用是什么?
  3. 下采样,是使用某种计算方式让更小的数据表示整张图片,这更小的数据代表了对这张图片高纬度的描述,而不是像素级细致的描述。
    1. 越小的数据对图片的表示和描述越宏观,有利于捕捉图片的语义特征。
    2. 例如一张猫在屋子前玩耍的地图,原图能看清所有细节,但因为细节太多,模型想要知道图里有猫和屋子,得把每个像素组合运算才行,但下采样到最小,最宏观的猫和屋子就容易识别。
  4. 上采样,是让图片的宏观小尺寸表示恢复成原图片尺寸。
    1. 比如对于图片分割(把图片上的物体分割出来),我们在下采样后的小数据量的高维表示里识别了图片的主体、边缘,最后还是要转回在原尺寸图片上表示,不然识别了也没用。
    2. 那不断下采样过程中肯定把图片细节都丢失了,再上采样,怎么可能还原图片细节?那就要说到跳跃连接(skip connection)了。
  5. 跳跃连接,也就是并不是顺着网络的方向连接,而是跳过原网络方向,跳着连接传输信息。说得有点拗口,看图很容易理解,就是图上中间的几条灰色箭头。
    1. 原网络连接方向是图片输入→下采样各节点→上采样各节点→输出图片这个链路,就是图中U字型的路径。
    2. 在这个路径之外,左边的下采样的每一层,都额外连接到右边上采样对应的层上面,将两个网络进行拼接。
    3. 上采样每一层,都在拼接了左边下采样对应层的数据后,再一起作为下一层上采样的输入。
    4. 为什么这样做,很容易理解,左边的每一层网络都保留了图片不同程度的细节,右边的每一层因为是上采样过来的,只有宏观信息,没有图片细节,那把左边图片细节信息拼接过去,右边这个网络宏观特征和微观细节都具备了,每一层都有不同程度的对图片的宏观语义理解和微观细节,就能做各种事情了,包括图片分割、语义生成图片。

UNet 网络大致思路是这样,这里面具体的卷积运算和公式,不看应该不影响对整体思路和作用的理解。

Stable Diffusion UNet 结构

最初的 DDPM(去噪扩散概率模型),和后来改良的 LDM(潜在扩散模型),对 UNet 网络逐步做了一些改造,以适合扩散模型图生成的过程,SD 是基于 LDM 实现的。

最后 SD 里的 UNet,整体结构流程跟上述一致,改造大部分是在上采样和下采样的每一层的实现里,最大的改造是引入了 ResnetBlock(残差模块)和 Transformer 模块。ResnetBlock 提升网络表达能力(原 UNet 是简单卷积模块),而 Transformer 模块的交叉注意力机制,将文本提示(prompt)的嵌入与图像特征进行融合,实现基于文本条件的图像生成。

SD UNet 每个模块的组成如图(图片来源):

2

左边下采样每层由2个残差模块和2个Transformer模块连接组成,右边上采样是各3个,中间层是2个残差模块和1个Transformer模块。(高维的d4和u1没有接入Transformer模块,原因不明,可能是试过加入后效果不佳,在高维这里加入 Prompt 交叉注意机制,文字权重太大?)

细分模块结构

里面每一块具体的结构这篇文章画得很详细,摘录学习一下。我们拿其中一个下采样模块看看:

3

两个残差模块,两个Transformer模块。这图表示了 SD 生图的三个输入:input(噪声图)、prompt_embdding(文字 Prompt)、time_embdding(步数)在这几个模块的流转和处理。这里每一个小模块处理完后,输出的可以近似认为都是一个预测的噪声图的数据表示。

残差模块的输入输出 噪声图+步数 → 预测噪声图,Transformer 模块的输入输出是 噪声图+ Prompt → 预测噪声图。

Transformer 模块

再细看一下 Transformer 模块,Transformer 模块由下图所示好几个部分组成,最主要的是 自注意力模块(SelfAttention)和交叉注意力模块(CrossAttention):

4

展开看看这两个模块:

5

自注意力模块,Transformer 结构里的 QKV 输入都是图片特征(上一层的处理结果,就是降噪图的特征),这样做可以让模型获得包含整个输入图像的感受野,捕捉图片特征中不同位置之间的关系, 全局感受力是 Transformer 架构的特点。

交叉注意力模块,它的作用是融合不同模态的输入,在这里就是融合噪声图和文本特征,Q的输入是图片特征,KV的输入是文字 prompt_embedding,让图片特征可以关注到文字输入,根据注意力权重调整图片的生成方向。文字 prompt 在整个Transformer模块中只作用在交叉注意力这部分里。

Transformer 的机制原理、QKV的含义,是另一个比较大的话题,可以先看看网上其他相关讲解,比如这篇,后续再细拆深入。

回顾

关键几个模块的组成了解了,再回到整个UNet的构成:

6

现在通过这些结构图,可以大致看到 UNet 网络里的整体处理流程,以及关键模块的作用,经过这些模块的逐个叠加,组合成一个个采样模块,再组合成 UNet 网络架构,完成整个生图运算。

这里面还有很多需要深入学习的点,当前先了解到这个维度,已经可以帮助大致理解后续 ControlNet 等网络的机制原理。

参考资料

UNet 论文:https://arxiv.org/abs/1505.04597

原版实现及 Diffusers 实现源码解读:https://zhouyifan.net/2024/01/23/20230713-SD3/周弈帆的博客,看简介挺有意思的一人)

U-Net简明教程:http://www.bimant.com/blog/unet-crash-tutorial/

Stable Diffusion1.5网络结构-超详细原创:https://blog.csdn.net/xd_wjc/article/details/134441396

Stable-Diffusion模型结构详解:https://zhuanlan.zhihu.com/p/638867353

Unet网络详解:https://blog.csdn.net/qq_58529413/article/details/125704059

Stable Diffusion XL网络结构:https://blog.csdn.net/xd_wjc/article/details/134530784

从零开始学扩散模型:https://huggingface.co/datasets/HuggingFace-CN-community/Diffusion-book-cn

分类:技术文章 Tags:
评论