Understand UNet in SD

Stable Diffusion 里的 UNet 主要有三个输入，分别是：

这三个输入经过 UNet 后，会得到一个输出，维度是 (1, 4, 64, 64)。

本文的目标就是：解释这个输出具体是怎么得到的。

整体结构

UNet 的结构主要包含五个部分：

• conv_in
• Encoder
• Middle Block
• Decoder
• conv_out

这里的 conv_in 和 conv_out 其实就是一个普通的 3x3 卷积层，用来转换通道数。

写成代码就是一行的事：

conv_in = Conv2d(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
conv_out = Conv2d(320, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

所以 UNet 的大部份操作，其实都是在 (1, 320, 64, 64) 这个维度基础上完成的。

先分析 Encoder

输入 (1, 320, 64, 64) 进入Encoder后，会依次经过四个Block：

1. CrossAttnDownBlock2D
2. CrossAttnDownBlock2D
3. CrossAttnDownBlock2D
4. DownBlock2D

得到输出 (1, 1280, 8, 8)。

可以看到，这里有两类 Block：

• CrossAttnDownBlock2D
• DownBlock2D

他俩的关系是：CrossAttnDownBlock2D 是加强版的 DownBlock2D，里面有 Attention Block ，而 DownBlock2D 里没有。

下面先介绍 DownBlock2D。

DownBlock2D 包含 N个ResBlock（N=2）和一个可选的下采样层。

特征会依次经过这N个ResBlock，最后经过下采样层（如果有）。

而 CrossAttnDownBlock2D 在每个ResBlock后，加了一个 Attention Block，用来计算图像特征的自注意力和图文特征的交叉注意力。

特征会依次经过这N对ResBlock-Attention Block，最后经过下采样层（如果有）。

对于Encoder里下采样层的设计是，前三个Block有，最后一个没有。

对于特征通道数的加倍是，中间两个会加倍，其他的不变。

下面是整体的设计：

conv_in、Encoder里的每个ResBlock以及下采样层的输出，都会存入down_block_res_samples，用于后续Decoder都跨层特征融合。

所以down_block_res_samples里一共有：1+2*4+1*3=1+8+3=12 个特征。

这12个特征的维度如下图：

再分析 Decoder

Encdoer 的输出会经过 Middle Block，并且特征维度和通道数都保持不变。

• 输入 x: (1, 1280, 8, 8)
• 内部层: 1 x ResBlock, 1 x Attention Block, 1 x ResBlock。
• 输出 x: (1, 1280, 8, 8)

Middle Block 的最终输出，也就是 Decoder 的输入，就需要和 Encdoer 的输出列表不断的进行拼接融合了。

输入在 Decoder 里也会经过四个 Block：

1. UpBlock2D
2. CrossAttnUpBlock2D
3. CrossAttnUpBlock2D
4. CrossAttnUpBlock2D

Decoder的设计和Encoder差不多，但是有以下区别：

• 前3个block要做上采样
• 后两个block做通道数减半
• 每个block包含3对ResBlock- Attention Block，而不是2对。

跨层分析

使用残差的时候，就是先拼接，再通过ResBlock融合特征通道。

因为前两个block对特征通道数不变，所以这三个ResBlock的输出维度是一样的。

up_blocks[0]

• 输入 sample_in: [B, 1280, 8, 8] (来自 mid_block)

这里也是，因为总体上前2个block的输出通道数是不变的，只是要根据残差的通道数来修改输入通道数。

up_blocks[1]

• 输入 sample_in: [B, 1280, 16, 16] (来自 up_blocks[0])

后两个Block对第一个ResBlock需要做到通道数减半。

up_blocks[2]

• 输入 sample_in: [B, 1280, 32, 32] (来自 up_blocks[1])

第一个ResBlock要做到通道数减半。

up_blocks[3]

• 输入 sample_in: [B, 640, 64, 64] (来自 up_blocks[2])