UNet结构拆解与PyTorch手敲：跳跃连接、上采样和形状对齐

UNet 这个网络做图像分割的人基本绕不开。我第一次接触它是在处理一批医学影像的语义分割任务时手头只有几百张标注图显存也不宽裕试过几个更“现代”的骨干网络效果反而没有一个结构朴素的小 UNet 稳。后来我养成了一个习惯凡是要长期用的网络先把结构从零手敲一遍不 import 任何现成的分割库。手敲的过程会逼着你去想清楚每一层的输入输出形状、通道数怎么翻倍、跳跃连接到底拼在哪一维上这些细节看别人的代码是看不太出来的。这篇内容就围绕 UNet 的网络结构拆解和代码手敲展开从编码器、解码器、跳跃连接的原理到 DoubleConv 的写法、上采样方式的选择、形状对齐的排查再到训练前的自检和常见坑全部按我实际写代码的顺序走一遍。适合刚入门语义分割、想搞清楚 UNet 内部到底在干什么的人也适合那些能跑通别人代码但一改结构就报错的同学。1. 为什么值得亲手把UNet敲一遍1.1 UNet到底解决了什么问题图像分割和分类的任务目标完全不是一回事。分类只需要告诉网络“这张图里有猫”输出一个概率向量就够了分割要求网络对每一个像素都给出判断输出和输入在空间上必须逐像素对应。这个差别决定了网络结构分类网络可以一路卷积加池化把空间尺寸压到 1×1分割网络不行它必须把压缩掉的空间信息再还原回来。UNet 在 2015 年被提出时最核心的贡献不是某个复杂模块而是给出了一个干净利落的对称结构左边一路下采样提取语义右边一路上采样恢复分辨率中间用跳跃连接把左边的浅层特征直接送到右边对应的层级。这个设计在当时解决了一个很实际的矛盾——深层特征语义强但位置模糊浅层特征位置精确但语义弱跳跃连接把两者拼在一起让解码器同时拿到“是什么”和“在哪里”。放到今天看这套思路依然是绝大多数分割网络的基本骨架。你去看近几年那些结构花哨的分割模型编码解码的对称框架基本没变变的只是卷积块内部换了什么、注意力加在哪、跳跃连接怎么融合。所以把 UNet 敲明白等于把分割网络的地基打好了。1.2 抄现成实现和手敲一遍的差别我见过太多人包括早期的我自己直接 pip 装个库一行调用把模型建出来训练脚本一跑指标还行就以为掌握了。等到要改输入通道数、要换损失函数、要在跳跃连接上加个门控、要把上采样从转置卷积换成双线性插值立刻就卡住因为根本不知道改哪一行会影响什么。手敲一遍的价值在于你会被迫回答几个具体问题第一层卷积的输入通道是几每个 stage 的通道数按什么规律翻倍下采样用的是最大池化还是步长为 2 的卷积上采样之后的特征图和跳跃连接送过来的特征图尺寸差一个像素怎么办拼接是在通道维还是空间维输出层的通道数等于类别数还是类别数减一这些问题单看论文的示意图是回答不了的示意图里那些方块并不标注 padding 和 stride。我个人的经验是第一次手敲大概率会在形状不匹配上卡两三个小时但卡完之后你对这个网络的理解会比读十遍论文都扎实。2. 先把结构拆开看编码器、解码器与跳跃连接2.1 编码器一路下采样的特征提取UNet 的编码器也叫收缩路径本质上就是一个标准的卷积分类网络只不过不接全连接层。原始论文里是 4 次下采样每次包含两个 3×3 卷积每个卷积后面跟 ReLU然后接一个 2×2 最大池化把空间尺寸减半。这里有个很容易被忽略的细节论文里用的是 valid 卷积也就是不加 padding所以每做一次 3×3 卷积特征图的长宽各减 2。这意味着输入图像不能随便给尺寸必须是 16 的倍数严格说是能被 2 的 4 次方整除并且要考虑每次卷积减 2 的累积效应论文里对输入做了镜像补边来规避这个问题。现在大家写代码基本都统一用 padding1让卷积保持尺寸不变只在池化时减半这样形状计算清爽很多。通道数的变化规律是 64 → 128 → 256 → 512 → 1024。每一级下采样之后通道翻倍、空间减半这是一个很自然的权衡空间信息在变少那就用更多的通道去承载更抽象的语义。我第一次看这个数字排布觉得是拍脑袋定的后来自己调过几组发现 64 起步、逐级翻倍在参数量和显存之间确实是个比较舒服的位置。起步通道给 32 会明显欠拟合给 128 则显存直接翻四倍小数据集上还容易过拟合。如果用一句话概括编码器的职责把一张 H×W×C 的图逐步压成 H/16×W/16×1024 的特征同时让每个位置上的向量包含足够大的感受野能“看懂”周围一大片区域的内容。2.2 解码器怎么把分辨率还回去解码器扩张路径要做的事和编码器相反把被压小的特征图一步步放大回原图尺寸同时把语义信息“摊开”到每个像素上。每一级包含一个上采样操作论文里是 2×2 的转置卷积把空间尺寸翻倍、通道数减半然后把编码器对应层级的特征拼过来再接两个 3×3 卷积做特征融合。这里有个关键点值得说清楚上采样本身不产生任何新信息它只是把已有的特征按某种规则铺开。转置卷积的可学习参数能让这个“铺开”的过程带一点自适应性但本质上仍然是插值。真正让解码器有能力做出精细判断的是后面拼接进来的跳跃连接特征。所以如果你把跳跃连接去掉只留一条编码-解码的主干模型是能训的但边界会糊得厉害小目标基本丢光。这一点我在一个细胞分割的数据集上验证过去掉跳跃连接之后 IoU 掉了将近 20 个点主要是边界和小面积目标崩了。解码器的最后通常还有一个 1×1 卷积把通道数压到类别数。这个 1×1 卷积的作用是把每个像素位置上的特征向量映射成一个分类得分可以理解成对每个像素独立做一次线性分类。输出通道数等于类别数多分类配合交叉熵或者等于 1二分类配合带 logits 的 BCE 损失。2.3 跳跃连接UNet真正的灵魂跳跃连接的具体做法是把编码器第 i 级输出的特征图和在解码器第 i 级上采样之后得到的特征图在通道维度上拼接concatenate而不是相加。这个选择很重要拼接保留了两种特征的独立性让后续的卷积自己去学怎么融合相加则强制要求两边通道数一致并且隐含地假设两种特征在同一语义空间里这在浅层和深层特征之间通常不成立。拼接带来的直接后果是通道数翻倍。假设解码器某级上采样后是 64 通道编码器对应层也是 64 通道拼完就是 128 通道所以后面的卷积块输入通道要按 128 来写。这是手敲代码时最容易写错的地方之一我建议在写每一级的时候把“上采样后通道数”“跳跃连接通道数”“拼接后通道数”三个数写在注释里能省掉大量调试时间。另外要注意层级对应关系。编码器有 5 个 stage4 次下采样解码器有 4 次上采样加上最底下那一层bottleneck一共 5 个层级。跳跃连接连的是编码器 stage1 到 stage4 的输出分别对应解码器的第 4 次到第 1 次上采样之后。最底下那一层通道 1024没有跳跃连接因为它就是编码器的最深处直接作为解码器的起点。2.4 通道数的排布规律与参数计算把整个网络的通道变化列成表会清楚很多。以标准的 4 次下采样版本为例假设输入是 1 通道灰度医学图像输出是 2 类层级操作输出形状H×W×C说明输入-572×572×1论文原尺寸实际常用 256 或 512enc1DoubleConv568×568×64padding0 时减 4pool1MaxPool2d(2)284×284×64尺寸减半enc2DoubleConv280×280×128通道翻倍pool2MaxPool2d(2)140×140×128-enc3DoubleConv136×136×256-pool3MaxPool2d(2)68×68×256-enc4DoubleConv64×64×512-pool4MaxPool2d(2)32×32×512-bottleneckDoubleConv28×28×1024无跳跃连接up4ConvTranspose2d56×56×512通道减半cat4concat with enc456×56×(512512)1024注意 enc4 需要裁剪dec4DoubleConv52×52×512-这张表是按原论文 valid 卷积写的你能看到 enc4 输出是 64×64而 up4 输出是 56×56两者对不上所以论文里专门做了一次中心裁剪。这也是为什么很多人复现原论文时形状老是报错。我实际写代码时一律用 padding1 的 same 卷积表格会干净得多层级操作输出形状备注输入-256×256×3-enc1DoubleConv(3,64)256×256×64-pool1MaxPool2d(2)128×128×64-enc2DoubleConv(64,128)128×128×128-pool2MaxPool2d(2)64×64×128-enc3DoubleConv(128,256)64×64×256-pool3MaxPool2d(2)32×32×256-enc4DoubleConv(256,512)32×32×512-pool4MaxPool2d(2)16×16×512-bottleneckDoubleConv(512,1024)16×16×1024-up4ConvTranspose2d(1024,512)32×32×512-cat4cat(up4, enc4)32×32×1024尺寸完全对齐dec4DoubleConv(1024,512)32×32×512-这样每一级的形状都能口算出来输入只要是 16 的倍数就不会出问题。3. 动手实现从卷积块到完整网络3.1 DoubleConv 双卷积块的写法UNet 里反复出现的单元是“两个 3×3 卷积 两个 ReLU”把它抽成一个类代码会干净很多。这里我用 PyTorch 写理由是调试时打印形状最方便print(x.shape)随处可插。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch, mid_chNone): super().__init__() mid_ch mid_ch or out_ch self.block nn.Sequential( nn.Conv2d(in_ch, mid_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_ch), nn.ReLU(inplaceTrue), nn.Conv2d(mid_ch, out_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)几个细节值得展开说一下。第一卷积后面我加了 BatchNorm原论文没有但现在的实践里基本都会加。BN 的作用是稳定每层输入的分布让网络对学习率不那么敏感在小 batch 上训练时收敛会快很多。代价是 batch size 太小比如 2时 BN 的统计量不准这时候可以换成 GroupNorm把 group 数设成 8 或 16在小 batch 的分割任务里更稳。第二卷积我设了biasFalse因为后面紧跟 BNBN 的 beta 参数本身就承担了偏置的作用再加一个 bias 是冗余的白白多占参数。这个习惯在写任何 Conv-BN 组合时都适用。第三inplaceTrue的 ReLU 能省一点显存代价是它会原地修改输入张量如果你的代码里有地方需要保留 ReLU 之前的特征就不能用 inplace。UNet 里没有这种需求所以放心用。第四我加了一个mid_ch参数默认等于out_ch这是为了后面做通道裁剪的变体留口子。标准的 UNet 用不到但如果你想压缩参数量把中间通道设成输出通道的一半参数能少将近一半精度掉得通常不多。3.2 下采样与上采样的实现选择下采样没什么争议nn.MaxPool2d(2)就够。选最大池化而不是平均池化理由是分割任务更关心“这个区域里有没有显著特征”最大值保留了最强的响应平均会把弱响应和强响应混在一起。也有人用 stride2 的卷积做下采样好处是参数可学习坏处是参数量和计算量都上去了在小数据集上不一定划算。我的默认选择是最大池化只有在数据量很大、想让网络自己学下采样方式时才换成卷积。上采样的选择就多了至少有三种常见做法差异不小方式实现优点缺点适用场景转置卷积nn.ConvTranspose2d(c_in, c_out, 2, stride2)可学习表达能力强容易产生棋盘格伪影数据量充足时最近邻/双线性插值 卷积F.interpolate(x, scale_factor2, modebilinear)后接 3×3 卷积无伪影稳定参数量略少表达稍弱医学图像、小数据集反池化记录最大池化位置后还原保留位置信息实现麻烦PyTorch 无现成层特定场景棋盘格伪影这个事值得单独说。转置卷积在 stride 大于 1 时输出像素由不同数量的输入像素贡献重叠不均匀训出来的结果上会出现规则网格状的明暗纹路。我在一次眼底血管分割的任务里就遇到过放大看预测图能明显看到格纹。换成双线性插值加卷积之后纹路消失指标还涨了一点。所以如果你在小数据集上训练我建议默认用插值方案转置卷积留作备选。代码上这样写class Up(nn.Module): def __init__(self, in_ch, out_ch, bilinearTrue): super().__init__() if bilinear: self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv DoubleConv(in_ch, out_ch, in_ch // 2) else: self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 self.up(x1) # 处理尺寸差一个像素的情况 diff_y x2.size(2) - x1.size(2) diff_x x2.size(3) - x1.size(3) x1 nn.functional.pad(x1, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) x torch.cat([x2, x1], dim1) return self.conv(x)这里的align_cornersTrue是个小坑。PyTorch 在插值时align_corners决定角点对齐方式设 True 时角点像素完全对齐设 False 时按比例缩放。对于 2 倍上采样两种设置在数值上有细微差别但对分割精度影响很小。真正需要注意的是训练和推理必须保持一致不要一处设 True 一处设 False否则预测结果会有偏移。nn.functional.pad那几行是防御性代码。如果你的输入尺寸是 16 的倍数理论上不会出现尺寸差但一旦用了非标准输入比如 250×250上采样和跳跃连接的尺寸就可能差一个像素。这几行代码能自动补边对齐比直接报错友好得多。3.3 拼接操作用的维度对齐torch.cat([x2, x1], dim1)里的dim1是通道维因为 PyTorch 的张量布局是(N, C, H, W)。这一点看起来简单但我见过不少人写成dim-1结果拼在了宽度上形状变成一个很宽的图后面的卷积还能跑因为卷积不关心宽高具体数值但输出完全错乱而且不报错属于最难查的那类 bug。拼接顺序也有讲究。torch.cat([x2, x1], dim1)和torch.cat([x1, x2], dim1)在数学上等价后面的卷积可以学到任意排列但为了和主流实现对齐、方便对照别人的权重我习惯把编码器送来的特征x2放在前面。这个不影响结果纯粹是代码可读性问题。还有一个隐含约束拼接要求两个张量除了通道维以外其他维度必须完全一致。如果 batch size 不一致不太可能、或者空间尺寸差一个像素就会直接报错。所以真正需要防的是空间尺寸这也是上一节那段 pad 代码存在的原因。3.4 完整网络组装把上面的模块拼起来完整的 UNet 长这样class UNet(nn.Module): def __init__(self, in_channels3, num_classes2, base_ch64, bilinearTrue): super().__init__() c base_ch self.inc DoubleConv(in_channels, c) self.down1 nn.Sequential(nn.MaxPool2d(2), DoubleConv(c, c * 2)) self.down2 nn.Sequential(nn.MaxPool2d(2), DoubleConv(c * 2, c * 4)) self.down3 nn.Sequential(nn.MaxPool2d(2), DoubleConv(c * 4, c * 8)) self.down4 nn.Sequential(nn.MaxPool2d(2), DoubleConv(c * 8, c * 16)) self.up1 Up(c * 16, c * 8, bilinear) self.up2 Up(c * 8, c * 4, bilinear) self.up3 Up(c * 4, c * 2, bilinear) self.up4 Up(c * 2, c, bilinear) self.outc nn.Conv2d(c, num_classes, kernel_size1) def forward(self, x): x1 self.inc(x) # c, H, W x2 self.down1(x1) # 2c, H/2, W/2 x3 self.down2(x2) # 4c, H/4, W/4 x4 self.down3(x3) # 8c, H/8, W/8 x5 self.down4(x4) # 16c, H/16, W/16 x self.up1(x5, x4) # 8c, H/8, W/8 x self.up2(x, x3) # 4c, H/4, W/4 x self.up3(x, x2) # 2c, H/2, W/2 x self.up4(x, x1) # c, H, W return self.outc(x) # num_classes, H, W我把下采样和紧随其后的 DoubleConv 打包成一个Sequential这样 forward 里每个down就是一级行数和层级对应得很清楚。注释里标了每一级的形状写的时候不标过两天回来看就得重新推一遍。base_ch参数化之后你可以一行切换小模型和大模型。base_ch64是原版参数量大约 3100 万输入 3 通道、输出 2 类。base_ch32参数量降到约 780 万显存占用大概四分之一适合小数据集和笔记本显卡。base_ch16只有 200 万参数能做快速原型验证但精度上限有限。输出层我用的是 1×1 卷积加num_classes通道不接 softmax。原因是 PyTorch 的CrossEntropyLoss内部已经包含了 log_softmax如果模型输出再接 softmax 会造成重复归一化训练直接崩掉。这个坑我踩过症状是 loss 一直卡在一个很高的值不下降梯度几乎为零。3.5 输入尺寸、通道数的验证写完网络第一件事不是拿去训练而是造一个假的输入跑一次前向model UNet(in_channels3, num_classes2, base_ch64) x torch.randn(2, 3, 256, 256) y model(x) print(y.shape) # 期望 torch.Size([2, 2, 256, 256])这一步能一次性验证三件事形状对不对、通道数对不对、有没有隐藏的维度 bug。我建议再补一组非 2 的幂的尺寸比如 250×250看看网络能不能扛住。如果报错说明 pad 那段逻辑没生效或者写错了如果能跑通但输出尺寸不是 250说明某处的卷积 padding 设错了。还可以顺手验证一下不同num_classes的情况。二分类时num_classes2输出通道 2如果要做单通道输出配合BCEWithLogitsLoss就把num_classes设成 1输出形状变成(N, 1, H, W)但这时候目标标签也要相应地变成浮点类型形状对齐成(N, 1, H, W)否则会报维度不匹配。4. 训练前的关键配置与形状验证4.1 损失函数与输出通道的匹配分割任务里损失函数和输出通道的配对关系我觉得比网络结构本身更容易出错因为报错信息往往不直观。整理成表会清楚一些任务类型输出通道目标标签形状与类型损失函数备注多分类含背景num_classes(N,H,W) int64取值 0 到 C-1CrossEntropyLoss最常用二分类单通道1(N,1,H,W) float取值 0/1BCEWithLogitsLoss注意标签要 float多标签像素可属多类num_classes(N,C,H,W) floatBCEWithLogitsLoss每类独立 sigmoid类别极不平衡num_classes同上 int64CrossEntropy weight 或 Dice前景占比低于 5% 时必调医学图像分割里前景占比经常很低比如一个 512×512 的图里病灶只占几百个像素占比不到 1%。这种情况下纯交叉熵会被背景主导模型学到的策略就是“全预测背景”准确率能到 99% 但毫无用处。常见的应对是给 CrossEntropyLoss 传weight参数把前景类的权重调高或者叠加 Dice Loss。Dice 直接优化重叠度对类别不平衡天然不敏感我的习惯是0.5 * CE 0.5 * Dice大部分场景下比单用一种稳。还有一点如果用CrossEntropyLoss目标标签必须是long类型int64形状是(N, H, W)不带通道维。我遇到过有人把 one-hot 编码的(N, C, H, W)直接喂进去报错信息是维度不匹配但很难一眼看出问题在哪。记住这个规则CE 要的是类别索引不是 one-hot。4.2 用一次前向传播做形状自检正式训练前我习惯写一段自检脚本把几个关键形状都打印出来跑通了再开训。这段脚本大概长这样def sanity_check(model, in_ch3, num_classes2, size256, devicecpu): model model.to(device) x torch.randn(2, in_ch, size, size).to(device) model.eval() with torch.no_grad(): y model(x) assert y.shape (2, num_classes, size, size), f输出形状异常: {y.shape} # 检查参数量 total sum(p.numel() for p in model.parameters()) trainable sum(p.numel() for p in model.parameters() if p.requires_grad) print(f总参数: {total/1e6:.2f}M, 可训练: {trainable/1e6:.2f}M) # 检查梯度是否能回传 model.train() y model(x) loss y.mean() loss.backward() no_grad [n for n, p in model.named_parameters() if p.grad is None] print(f无梯度参数数量: {len(no_grad)}) if no_grad: print(前几个:, no_grad[:5])这段脚本的第二部分很关键。有时候网络形状对、能前向但某些分支的梯度断了训起来就是不动。常见原因是某个模块用了detach()、或者跳跃连接被写成了一个纯 tensor 拼接而没参与计算图、或者torch.no_grad()误加在了训练路径里。跑一次 backward 然后检查哪些参数没有梯度能快速定位。4.3 参数量与显存估算参数量手算其实不难UNet 的绝大部分参数集中在卷积层。以base_ch64、输入 3 通道为例inc那一层的第一个卷积是 3×3×3×641728 个参数第二个是 3×3×64×6436864加上 BN 的 128 个这一层不到 4 万。真正的大头在 bottleneckdown4的第二个卷积是 3×3×512×1024≈470 万bottleneck 内部两个卷积加起来接近 1200 万占了总参数量的三分之一还多。显存占用则主要取决于特征图的尺寸和 batch size而不是参数量。一张 256×256 的图inc输出的特征图是 256×256×64float32 下就是 16MB反向传播要存中间激活实际占用还要乘好几倍。这就是为什么有人发现模型只有 3000 万参数但 batch size 开到 8 就爆显存。几个实用的省显存手段把base_ch从 64 降到 32特征图通道减半激活占用直接降到四分之一用混合精度训练激活存成 float16显存再省近一半把输入从 512 裁到 256激活占用降到四分之一。我一般优先调输入尺寸因为它对精度的影响最可预期其次才是通道数。5. 常见问题与排查实录5.1 形状不匹配类问题手敲 UNet 遇到的报错八成都和形状有关。最典型的是拼接时报错RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 64 but got size 63 for tensor number 1 in the list.这个报错的意思是除了通道维其他维对不上。原因通常是输入尺寸不是 16 的倍数。比如输入 250×250经过四次池化250→125→62→31→15这里 125 是奇数池化时向下取整变成 62再往上采样就是 124和编码器的 125 差一个像素拼接直接炸。解决办法有两个一是严格要求输入尺寸是 16 的倍数在数据加载时统一 resize 或 pad 到 256/512二是在Up模块里加前面那段F.pad逻辑自动补边对齐。我两个都做双保险。第二个常见报错是输出通道数和标签类别数不匹配报错信息是Expected target size之类。这种一般是num_classes设错了或者标签里出现了超出范围的类别索引。有一次我的标签里混进了一个值为 255 的“忽略像素”而num_classes2CE 直接报索引越界。处理方法是在构造标签时把忽略像素映射到合法范围或者给 CE 传ignore_index255。5.2 上采样方式的坑转置卷积的棋盘格伪影前面提过了这里补充一个更隐蔽的问题ConvTranspose2d的output_padding参数。当输入尺寸是奇数时转置卷积的输出尺寸公式会有歧义需要靠output_padding指定多出来的那一行/列加在哪边。不设的话默认是 0输出尺寸可能比你期望的少 1然后拼接就报错了。这个参数我建议在写的时候就显式加上output_padding1能避免很多奇数尺寸下的诡异问题。另一个坑是nn.Upsample和F.interpolate的行为差异。两者本质一样但nn.Upsample是一个模块可以放进SequentialF.interpolate是函数式调用。如果用了nn.Upsample但没设align_cornersPyTorch 会给一个警告训练时刷屏。直接显式写上就没问题。还有个小细节双线性插值的modebilinear只支持 4D 输入N,C,H,W如果你的数据是 5D比如视频或 3D 医学图像要用trilinear。UNet 的 3D 版本在医学影像里很常用改的时候记得同步改这个。5.3 训练不收敛、预测全黑全白模型能跑通、loss 也在动但预测结果全是一类这个问题比形状报错更难查。我按排查顺序列一下我自己的检查清单第一步看数据。把标签可视化出来确认前景真的存在、标注没有全零。我遇到过一次数据加载时的 resize 用了最近邻标签里的类别索引被插值成了中间值比如 1 和 0 之间插出 0.5然后转成 long 变成 0前景全没了。标签的 resize 必须用NEAREST不能用默认的双线性。第二步看学习率。UNet 从零训学习率 1e-3 配 Adam 一般没问题如果用了预训练编码器学习率要降到 1e-4 甚至更低否则预训练权重会被冲掉。反过来学习率给太小1e-5会导致 loss 几乎不动看起来也像不收敛。第三步看损失函数和输出是否配套前面说过 softmax 重复的问题。还有一个隐蔽的如果模型输出接了 sigmoid 但用的是BCEWithLogitsLoss等于做了两次 sigmoid输出全被压到 0.5 附近看起来就是一片灰。第四步看 BN 的统计量。如果训练时 batch size 很小比如 2BN 的 running mean/var 估计不准推理时表现会明显差于训练时。判断方法是切换model.eval()和model.train()各跑一次如果指标差很多基本就是 BN 的问题。换成 GroupNorm 通常能解决。5.4 常见问题速查表把上面这些整理成一张表出问题的时候可以直接对现象最可能原因快速验证处理拼接时尺寸不匹配输入非 16 倍数或 padding 设错打印各层特征图尺寸输入 resize 到 16 倍数Up 模块加 pad输出通道与标签不符num_classes 设错或标签越界打印 y.shape 和标签 max检查类别数处理忽略像素loss 卡住不降输出重复 softmax或学习率过大检查模型最后一层去掉 softmax调小 lr预测全黑/全白类别极不平衡或标签 resize 用错插值统计前景像素占比加类别权重或 Dice标签用 NEAREST训练好推理差BN 统计量不准对比 train/eval 模式指标换 GroupNorm 或加大 batch预测有网格纹路转置卷积棋盘格伪影放大预测图看纹理换成插值上采样梯度为 None某分支被 detach 或未参与计算图反向传播后检查 p.grad移除 detach检查拼接是否在计算图内显存不够输入太大或通道太多打印单步显存占用降 base_ch、降分辨率、混合精度6. 我踩过的坑和一些改写建议6.1 几种常见的UNet改法把标准结构敲通之后理解变体会容易很多因为它们大多数只动一两个地方。第一种是通道裁剪版。把每级的通道数统一砍半32/64/128/256/512参数量降到原来的四分之一左右。我在一个只有 800 张标注图的任务上用过精度只掉了不到 1 个点但训练速度快了三倍多。小数据集上大模型未必更好这条在分割任务里体现得特别明显。第二种是残差版。把DoubleConv里的两个卷积改成带残差连接的形式缓解深层梯度消失。UNet 最深也就 5 级梯度消失不算严重但如果你把它加深到 6 到 7 级残差就很有必要了。注意残差要求输入输出通道一致通道变化的那一层需要用 1×1 卷积做投影。第三种是注意力门控版。在跳跃连接上加一个小的注意力模块让解码器有选择地接收编码器的特征抑制无关区域。这个改法在医学图像里用得很多因为病灶区域通常很小背景干扰大。实现上是在拼接之前用解码器的特征生成一个权重图乘到编码器特征上。参数量增加很少但要注意别加太深否则小数据集上容易过拟合。第四种是深度可分离卷积版。把标准 3×3 卷积换成深度可分离卷积参数量和计算量都大幅下降适合移动端部署。代价是精度通常会掉一点而且在小通道数上优势不明显深度卷积的通道数太少时效率不高。6.2 工程上的一些细节最后分享几个和结构无关但很影响实际使用体验的点。权重初始化。PyTorch 默认用 Kaiming 初始化对 ReLU 网络是合适的一般不用改。但如果你把激活换成了 SiLU 或 GELU最好手动指定一下初始化方式默认的可能会让训练初期不稳定。我遇到过一次换激活后 loss 前几十个 step 剧烈震荡换回默认初始化就正常了。输入归一化。UNet 的编码器没有全局池化对输入的尺度比较敏感。图像像素值直接除以 255 是最低要求更好的做法是按数据集统计均值和方差做标准化。医学图像里这个差别很明显因为不同设备的灰度分布差异很大。推理时的尺寸处理。训练用 256×256 的 patch推理时如果整张图是 1024×1024直接送进去会因为显存不够而失败或者因为尺寸不是 16 的倍数而报错。常见的处理是滑窗切 patch重叠区域取平均。切的时候注意 patch 的边界要留出足够的重叠我一般用 1/4 的重叠否则拼接处会有明显的接缝。保存和加载。保存state_dict而不是整个模型对象这样加载时不依赖原始类的定义路径。但这就意味着加载时必须先实例化同样的结构所以建议把模型的关键参数in_channels、num_classes、base_ch一起存进 checkpoint加载时读出来再建模型。我就吃过亏改了一次base_ch之后老权重怎么都加载不上查了半天才发现是通道数对不上。多卡训练时的 BN。如果用 DataParallel 或 DistributedDataParallel每个卡上的 batch 会被切分BN 统计量是按单卡算的。这时候要么把 BN 换成 SyncBN要么保证单卡 batch 足够大。这个细节在单卡调试时完全看不出来一上多卡就露馅。我个人在实际操作中的体会是手敲 UNet 最大的收获不是这个网络本身而是建立了一套“形状可推、参数可算、问题可查”的思维习惯。后面再去看 DeepLab、SegFormer 这些结构会发现它们的改动都能对应到具体的形状变化和参数取舍上不再是一堆看不懂的模块堆叠。如果你也想练手建议从一个 3 通道 256×256 的小数据集开始先把标准版跑通再试着改一个地方看看影响比一次性堆一堆改动的学习效率高得多。