""" unet.py —— U-Net：slides 的原始结构（Ronneberger et al., 2015）和一个给扩散模型用的小 U-Net（第 7–8 讲） ===================================================================== 【这个文件要回答的问题】 扩散模型的去噪网络要输入一张图、输出一张**同样大小**的图（预测的噪声 ε）。分类用的 CNN 一层层缩小、 最后只剩一个向量，不合适；U-Net 先缩小（收缩路径，看全局）、再放大（扩张路径，恢复分辨率）， 并把收缩路径每一层的特征直接"跳接"到扩张路径的同一层（拼接），让细节不在缩小的过程中丢掉。 【slides 的 U-Net（为语义分割设计）】 收缩路径的每个块：3×3 卷积 → 3×3 卷积 → ReLU → 步长 2 的最大池化（下采样）；重复 N 次，每次通道数翻倍。 扩张路径的每个块：2×2 卷积（上采样）→ 和收缩路径的特征拼接 → 3×3 卷积 → 3×3 卷积 → ReLU；重复 N 次，通道数减半。 关键：输出层和输入图像的长宽相同（通道数可以不同）。 （原论文每个 3×3 卷积后面都有 ReLU；卷积不补零，所以 572×572 的输入得到 388×388 的输出。） 【本文件提供】 OriginalUNet 原论文的结构（不补零的卷积、64→1024 通道），只用来数参数、推算形状（01 号脚本用 meta 设备，不分配内存） TinyUNet 给 28×28 形状图用的小 U-Net：先"像素重排"到 14×14、补零的卷积（输出和输入一样大）、 GroupNorm、SiLU、时间步嵌入加到每个块里；skip=False 时去掉跳接（对照实验用） 【边界】TinyUNet 要求输入的长宽能被 4 整除（像素重排一次、下采样一次）；28×28 → 14×14 → 7×7。 """ import torch import torch.nn as nn import torch.nn.functional as F from diffusion_lab import time_embedding # ===================================================================== # 1. 原论文的 U-Net（只用来数参数和推算形状） # ===================================================================== class _DoubleConv(nn.Module): """两个不补零的 3×3 卷积，各接一个 ReLU（每个卷积让长宽各减 2）。""" def __init__(self, cin, cout): super().__init__() self.net = nn.Sequential(nn.Conv2d(cin, cout, 3), nn.ReLU(), nn.Conv2d(cout, cout, 3), nn.ReLU()) def forward(self, x): return self.net(x) class OriginalUNet(nn.Module): """ Ronneberger et al. (2015) 的 U-Net：通道 64 → 128 → 256 → 512 → 1024，再对称地回来；输出 n_classes 个通道。 【跳接时的裁剪】不补零的卷积让收缩路径的特征图比扩张路径的大，拼接前要把它从中间裁成同样大小（论文图 1 的 "copy and crop"）。 """ def __init__(self, in_ch=1, n_classes=2, base=64, depth=4): super().__init__() chs = [base * 2 ** i for i in range(depth + 1)] # 64, 128, 256, 512, 1024 self.downs = nn.ModuleList() c = in_ch for ch in chs[:-1]: self.downs.append(_DoubleConv(c, ch)) c = ch self.bottom = _DoubleConv(chs[-2], chs[-1]) self.ups = nn.ModuleList() self.up_convs = nn.ModuleList() for ch in reversed(chs[:-1]): self.ups.append(nn.ConvTranspose2d(ch * 2, ch, 2, stride=2)) # "2×2 卷积（上采样）"，通道减半 self.up_convs.append(_DoubleConv(ch * 2, ch)) # 拼接后通道又翻倍，两个 3×3 卷积再减半 self.head = nn.Conv2d(chs[0], n_classes, 1) @staticmethod def _crop(x, target): dh = (x.shape[2] - target.shape[2]) // 2 dw = (x.shape[3] - target.shape[3]) // 2 return x[:, :, dh:dh + target.shape[2], dw:dw + target.shape[3]] def forward(self, x, trace=None): skips = [] for blk in self.downs: x = blk(x) skips.append(x) if trace is not None: trace.append(("收缩", tuple(x.shape[1:]))) x = F.max_pool2d(x, 2) x = self.bottom(x) if trace is not None: trace.append(("最底层", tuple(x.shape[1:]))) for up, conv, skip in zip(self.ups, self.up_convs, reversed(skips)): x = up(x) x = conv(torch.cat([self._crop(skip, x), x], dim=1)) if trace is not None: trace.append(("扩张", tuple(x.shape[1:]))) return self.head(x) # ===================================================================== # 2. 扩散模型用的小 U-Net # ===================================================================== class _Block(nn.Module): """ 补零的"两个 3×3 卷积"块，加上时间步： h = Conv3×3(SiLU(GroupNorm(x)))；h = h + Linear(时间嵌入)（每个通道加一个偏移）； h = Conv3×3(SiLU(GroupNorm(h)))；输出 h + 跳过这个块的捷径（残差，通道数不同时用 1×1 卷积对齐） 【为什么要把时间步加进每个块】同一张带噪的图，在 t = 50 和 t = 800 时该去掉的噪声量完全不同； 网络在每一层都需要知道"现在噪声有多大"。 【GroupNorm】把通道分成几组、在组内做归一化；和 BatchNorm 不同，它不依赖同一批的其他样本， 训练和采样时的行为一样（扩散模型采样时一批里的图噪声水平都相同，BatchNorm 的统计量会出问题）。 """ def __init__(self, cin, cout, t_dim, groups=8): super().__init__() self.n1 = nn.GroupNorm(min(groups, cin), cin) self.c1 = nn.Conv2d(cin, cout, 3, padding=1) self.t = nn.Linear(t_dim, cout) self.n2 = nn.GroupNorm(min(groups, cout), cout) self.c2 = nn.Conv2d(cout, cout, 3, padding=1) self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity() def forward(self, x, temb): h = self.c1(F.silu(self.n1(x))) h = h + self.t(temb)[:, :, None, None] h = self.c2(F.silu(self.n2(h))) return h + self.skip(x) class TinyUNet(nn.Module): """ 28×28 单通道图像的去噪 U-Net（c = 32 时约 31 万参数，第 8 讲用的 c = 24 约 19 万）： 输入 (1, 28, 28) → 像素重排 → (4, 14, 14) → 3×3 卷积 → (c, 14, 14) → 块 ──────────────┐跳接 → 步长 2 卷积 → (2c, 7, 7) → 块 → 块（最底层） │ → 转置卷积 → (c, 14, 14) → 和跳接拼接 (2c, 14, 14) → 块 → (c, 14, 14) ←───────────┘ → GroupNorm → SiLU → 3×3 卷积 → (4, 14, 14) → 像素重排的逆 → 输出 (1, 28, 28) 【像素重排 pixel unshuffle】把每个 2×2 的像素块"叠"成 4 个通道：(1, 28, 28) → (4, 14, 14)，一个数都不丢； 输出时再"摊"回去（pixel shuffle）。这样所有卷积都在 14×14 和 7×7 上做，计算量约为直接在 28×28 上做的四分之一。 我们的 CPU 上卷积很慢（直接在 28×28 上用 16 个通道的版本每步要 173 ms，这个版本约 90 ms），这是为了让 第 8 讲的训练能在几分钟内完成而做的妥协。DiT、U-ViT 等扩散 Transformer 的"切块"（patchify）是同一个思路（第 14 讲）。 和 slides 的其他差别：下采样用步长 2 的卷积代替最大池化、上采样用 4×4 的转置卷积，卷积都补零（输出和输入一样大）， 每个块是残差块并且加入了时间步——这些都是 DDPM 及之后扩散模型 U-Net 的常见做法。U 形的结构 （收缩 → 扩张 + 同一分辨率的跳接拼接）和 slides 相同。 【参数】c：14×14 那一层的通道数；t_dim：时间嵌入维数；skip：是否用跳接（False 时扩张路径不拼接，对照实验用） """ def __init__(self, c=32, t_dim=64, skip=True): super().__init__() self.t_dim, self.use_skip = t_dim, skip self.t_mlp = nn.Sequential(nn.Linear(t_dim, 4 * t_dim), nn.SiLU(), nn.Linear(4 * t_dim, t_dim)) self.inc = nn.Conv2d(4, c, 3, padding=1) self.d1 = _Block(c, c, t_dim) self.down = nn.Conv2d(c, 2 * c, 4, stride=2, padding=1) self.mid1 = _Block(2 * c, 2 * c, t_dim) self.mid2 = _Block(2 * c, 2 * c, t_dim) self.up = nn.ConvTranspose2d(2 * c, c, 4, stride=2, padding=1) self.u1 = _Block(2 * c if skip else c, c, t_dim) self.out = nn.Sequential(nn.GroupNorm(8, c), nn.SiLU(), nn.Conv2d(c, 4, 3, padding=1)) def forward(self, x, t, trace=None): temb = self.t_mlp(time_embedding(t, self.t_dim)) h0 = F.pixel_unshuffle(x, 2) # (4, 14, 14) h1 = self.d1(self.inc(h0), temb) # (c, 14, 14) h = self.mid2(self.mid1(self.down(h1), temb), temb) # (2c, 7, 7) h = self.up(h) # (c, 14, 14) h = self.u1(torch.cat([h, h1], 1) if self.use_skip else h, temb) out = F.pixel_shuffle(self.out(h), 2) # (1, 28, 28) if trace is not None: trace += [("像素重排", tuple(h0.shape[1:])), ("收缩：块", tuple(h1.shape[1:])), ("最底层", (2 * h1.shape[1], h1.shape[2] // 2, h1.shape[3] // 2)), ("扩张：上采样 + 拼接 + 块", tuple(h.shape[1:])), ("输出", tuple(out.shape[1:]))] return out