Diffusers中UNet2D上采样块ResNet层数与通道数疑问

问题背景 #

我正在学习Unet与Stable Diffusion，计划训练模型，已理解原始Unet架构及下采样、上采样块的通道、尺寸变化逻辑。我的理解是：若上采样块含2个ResNet块，仅第一个ResNet块的首个Conv2d操作因接收前序模块+跳跃连接输入而通道数翻倍，其余ResNet块输入通道与前一块输出通道一致，此理解是否正确？

在使用Diffusers库实现以下UNet2DModel时，对其上采样块的ResNet块工作逻辑存在困惑：

from diffusers import UNet2DModel
# 训练64×64尺寸的图像
image_size = 64
# 创建模型
model = UNet2DModel(
    sample_size=image_size,  # 目标图像分辨率
    in_channels=4,  # 输入通道数， latent图像为4，像素空间为3
    out_channels=4,  # 输出通道数
    layers_per_block=2,  # 每个UNet块包含的ResNet层数
    block_out_channels=(64, 128, 256, 512), 
    down_block_types=(
        "DownBlock2D",  # 常规ResNet下采样块
        "DownBlock2D",
        "AttnDownBlock2D",  # 带空间自注意力的ResNet下采样块
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",  # 带空间自注意力的ResNet上采样块
        "UpBlock2D",
        "UpBlock2D",  # 常规ResNet上采样块
    ),
)

核心疑问 #

1. 为什么layers_per_block设为2时，下采样块各含2个ResNet块，但上采样块却有3个？
2. 首个AttnUpBlock2D的3个ResNet块中，第二个ResNet块输入通道为1024、第三个为768的计算逻辑是什么？已知输入latent为64×64×4，中间块输出8×8×512，最后一个下采样块跳跃连接也是8×8×512，仅能理解第一个ResNet块首个Conv2d的1024通道来源。

解答 #

初始理解验证 #

你的初始理解是正确的：当采样块包含2个ResNet块时，只有第一个ResNet块的首个Conv2d会因融合前序模块输出与跳跃连接特征，导致输入通道数翻倍；后续ResNet块的输入通道数与前一个ResNet块的输出通道数完全一致。

疑问1解答 #

Diffusers库的UNet2DModel中，上采样块与下采样块的结构设计存在差异：

• 下采样块的结构为：layers_per_block个ResNet层 + 1个下采样层（非最后一个下采样块时），因此ResNet层数量完全等于layers_per_block的设置值。
• 上采样块的结构为：1个上采样层 + layers_per_block个ResNet层 + 1个额外ResNet层（专门用于处理融合跳跃连接后的特征）。当layers_per_block=2时，上采样块的ResNet层总数就是2+1=3个。

疑问2解答 #

我们以首个AttnUpBlock2D为例拆解通道计算逻辑：

1. 第一个ResNet块输入1024：来自中间瓶颈块输出（8×8×512）与对应下采样块的跳跃连接特征（8×8×512）的拼接，512+512=1024，这部分你已经理解。
2. 第二个ResNet块输入1024：ResNet块的默认设计是保持输入输出通道数一致（除非有专门的通道调整卷积或下采样/上采样操作），因此第一个ResNet块的输出通道数仍为1024，直接作为第二个ResNet块的输入。
3. 第三个ResNet块输入768：第二个ResNet块输出后，模型会通过一个1×1卷积层将通道数从1024调整为768。这个数值是模型为了匹配后续上采样块的输入要求设计的——768=512（当前块的基础通道）+256（下一个上采样块的目标通道），完成通道过渡后，再送入第三个ResNet块处理。

内容的提问来源于stack exchange，提问作者Krishna Dave