第一次把 Unet 跑起来的时候最让人抓狂的往往不是 loss 不下降而是训练到一半突然蹦出一句张量形状对不上的报错。很多人对 Unet 的第一印象停留在一个 U 形的编码器加解码器加上几根跳跃连接看起来结构简单到几乎朴素。但真正动手复现过、在真实数据集上训过、调过 loss 和 batch 的人都知道这个 2015 年提出的 Unet 网络结构图背后藏着一堆需要自己想明白的细节为什么必须拼特征而不是加特征、为什么医学图像上它比同期全卷积网络好用、为什么输入尺寸非要卡在 16 的倍数、类别极端不平衡时该换哪个损失函数。这篇就把 Unet 网络详解从结构图拆到每一层的张量推导再讲到训练时的坑和改进方向尽量把那些文档里不会写、但自己踩过才知道的东西讲清楚。1. 从一张 512×512 的医学影像说起Unet 到底在解决什么麻烦1.1 语义分割和图像分类根本不是一类任务图像分类只要求你回答这张图里有什么模型只要在最后输出一个类别概率就够了中间过程怎么压缩、怎么抽象都无所谓。但 Unet 图像分割要的是逐像素的答案——每个像素属于前景还是背景、属于哪一类细胞、属于哪块肿瘤。这意味着网络不能一路把空间信息压没了它必须在降维理解语义的同时把这个像素在哪这件事守恒下来。这就像翻译分类任务是把一整段话翻译成一个标题只要意思对就行分割任务是把一段话逐词逐句翻译位置、顺序、细节一个都不能错位。所以任何只为分类设计的骨干网络直接拿来做分割都会有一个通病——深层特征语义很强但空间分辨率太低上采样回来的边界是糊的。Unet 的全部设计动机就是冲着这个矛盾去的。1.2 2015 年之前的分割网络卡在哪里在 Unet 之前全卷积网络Fully Convolutional Network已经证明了可以用卷积替代全连接层来实现逐像素预测思路是把分类网络的后半段改成上采样。问题是它的上采样依赖很浅的高层特征配一个比较大的步长直接放大回去细节丢得厉害。医学图像又特别吃亏——细胞边界、细小血管、病灶边缘这些恰恰是最需要精确的地方一旦糊掉临床价值就没了。还有一个现实约束当年医学影像数据集普遍很小ISBI 细胞追踪挑战赛那种级别能用来训练的标注图像也就几十张。参数量动辄上亿的大网络根本喂不饱过拟合会非常严重。Unet 之所以能在那个比赛里拿到好成绩一部分原因就是它的结构本身就带很强的先验——它假设不同尺度上的信息都有用这个假设在生物医学图像上几乎总是成立。1.3 Unet 给出的答案把结构当成一种先验Unet 的核心主张很朴素既然深层有语义、浅层有位置那就把浅层的位置信息直接搬到深层去用。做法是让编码器和解码器在每一级分辨率上对齐解码器上采样到某个尺寸时就把编码器同尺寸的特征图接过来一起用。这样解码器在恢复分辨率的时候手里既有抽象语义又有相对干净的空间轮廓。这个设计的美妙之处在于它没有引入任何新的算子全部由卷积、池化、上采样和拼接组成普通显卡都跑得动。也正因为如此后来几乎所有分割网络的改进都是在 Unet 这个骨架上做加法——换骨干、加注意力、改成多尺度输出而那一句编码器负责语义、解码器负责定位、跳跃连接负责把它们缝合起来成了整个分割领域的通用认知。2. 把 U 型结构拆到每一层编码器、瓶颈层、解码器各自在做什么2.1 编码器四次下采样换来的语义抽象编码器这一侧做的事可以用一句话概括——把空间换成语义。原始 Unet 里每一级都由两个 3×3 卷积加 ReLU 组成然后接一个 2×2 最大池化把尺寸砍半、通道翻倍。通道数的走法是 64、128、256、512、1024一共四次池化宽高总共缩小 16 倍。为什么非要重复两个卷积而不是一个因为两个 3×3 叠起来的感受野等于一个 5×5但参数量和计算量更小同时中间多了一层非线性表达能力更强。而且第一个卷积负责从上一级特征里提取新响应第二个卷积负责在这个响应上再组合层与层之间的分工更清晰。这个两层卷积一池化的块结构后来被反复沿用成了分割网络的默认积木。池化的作用不只是降维。它带有一点平移不变性让网络对小范围位移不那么敏感代价是分辨率下降。所以池化是以精度换语义的买卖Unet 选择在下采样阶段做这个买卖然后靠解码器和跳跃连接把精度买回来。2.2 瓶颈层特征图最小、语义最浓的那一层经过四次下采样之后特征图的宽高只有输入的 1/16通道数达到 1024这一层通常叫瓶颈层。它的空间维度最小但每个像素背后对应的感受野已经覆盖了原图很大一块区域所以它表达的是全局语义——整块组织、整片区域是什么而不是某个局部细节。很多人在讲 Unet 网络结构图时会跳过瓶颈层觉得它只是中间的一环。实际上它是整个网络信息流的拐点所有往下的压缩在这里到头接下来全是往上的恢复。如果你在这一层插入注意力、全局池化或者 Transformer 模块收益通常比在浅层加更明显因为这里的特征最浓缩改动一点点就能影响全局。但瓶颈层也很脆弱。它的空间分辨率极低如果上采样路径设计得不好错误会在这里被放大。所以解码器与瓶颈层之间必须保证通道数和尺寸的衔接是干净的不能出现从 1024 通道直接跳到 128 通道这种断崖式变化。2.3 解码器上采样怎么把语义重新翻译回像素解码器一侧由三部分组成上采样、拼接、两层卷积。上采样原始论文用的是 2×2 转置卷积把尺寸翻倍、通道减半然后与编码器同尺寸的特征图拼接通道数变成两倍最后用两层 3×3 卷积把通道数压回去。每经过这样一次操作宽高翻倍通道数减半一路恢复到接近原图的尺寸。转置卷积有个众所周知的毛病——当卷积核尺寸和步长不能整除时输出会出现棋盘格状伪影。现在很多实现干脆改成最近邻或双线性插值上采样 一个 3×3 卷积效果更稳也不容易出奇怪的花纹。我个人的经验是用插值 卷积的组合在小数据集上收敛更平顺转置卷积偶尔会出现某些位置被过度激活的情况需要额外调初始化。2.4 一次完整的张量形状推演我们把原始 Unet 的尺寸变化走一遍。原始论文为了做无重叠的 valid 卷积输入是 572×572每经过一次 3×3 卷积宽高减 2所以每个卷积块内两层卷积共减 4。具体走法如下阶段操作输出尺寸宽高输出通道输入-572×5721编码块12×3×3 conv568×56864池化12×2 maxpool284×28464编码块22×3×3 conv280×280128池化22×2 maxpool140×140128编码块32×3×3 conv136×136256池化32×2 maxpool68×68256编码块42×3×3 conv64×64512池化42×2 maxpool32×32512瓶颈2×3×3 conv28×281024从瓶颈往上走上采样把 28×28 变成 56×56可编码器对应那一级经过卷积后是 64×64尺寸并不完全相等。这时候原始实现是先把编码器特征裁剪到 56×56 再拼接所以论文里输入输出尺寸差异很大——输入 572输出只有 388。这个裁剪逻辑就是当年很多人复现失败的元凶之一。提示现在主流的 PyTorch、TensorFlow 实现基本都用 same padding 的 3×3 卷积尺寸在每层保持不变拼接时不需要裁剪输入输出尺寸也就一致了。如果你照着原始论文的数字去写代码反而会处处对不上。3. 跳跃连接才是 Unet 的灵魂它到底接回了什么3.1 没有跳跃连接会发生什么如果把跳跃连接全部去掉Unet 就退化成一个普通的编码器-解码器网络也就是自编码器做分割。这种网络训出来的结果区域大轮廓往往还行但边界会明显偏软、偏胖或偏瘦小目标容易直接消失。原因是解码器上采样时手里的信息全部来自深层深层的分辨率已经很低每个像素代表原图一大块区域靠它恢复出的边界天然是平均化的。跳跃连接接回的那部分特征恰恰是还没被池化严重压缩的浅层特征。它们语义弱但对边界的响应强因为浅层卷积核本质上就是在检测边缘、角点、纹理。把这两类信息拼在一起解码器才同时拥有这是什么东西和它的边在哪里两种知识。3.2 拼接和相加的区别为什么 Unet 选拼接特征融合有两大做法相加add和拼接concat。相加是把两张同尺寸同通道的特征图直接逐元素相加通道数不变拼接是把它们在通道维度上叠起来通道数翻倍。ResNet 那种残差连接用的是相加Unet 用的是拼接。拼接的好处是信息不做混合浅层特征和深层特征各自保留后面的卷积可以自己学着去挑选有用的通道。相加相当于强行假设两边通道一一对应且语义可线性叠加这个假设在分割任务里并不成立——浅层的边缘通道和深层的语义通道根本不是一类东西硬加会互相干扰。代价也很直接拼接后通道数翻倍紧接着的卷积和显存开销也随之翻倍这是 Unet 训练吃显存的主要原因之一。3.3 拼接带来的显存代价与通道爆炸以第一级解码器为例上采样得到 512 通道、瓶颈上一级尺寸的特征与编码器同级的 512 通道特征拼接瞬间变成 1024 通道随后要接两层 3×3 卷积。这两层卷积的参数量是 1024×512×9 和 512×512×9加起来超过 700 万比同级编码器还大。整个网络的参数量大约在 3100 万左右其中相当一部分就花在这些拼接之后的卷积上。所以如果你的显存吃紧一个立竿见影的优化点是减少拼接后的卷积通道或者在拼接前做一次 1×1 卷积降维。这种做法会略微损失精度但对推理速度和显存占用帮助明显。另一个思路是把最浅层的拼接去掉——最浅层特征图尺寸最大、通道虽然少但空间位置多显存占用其实相当可观。4. 参数量与感受野几组必须自己算一遍的数4.1 逐层估算参数量和显存占用很多人对卷积参数量没有直观概念觉得 3×3 卷积核小应该很省。实际上参数量取决于输入通道乘输出通道通道一多就爆炸。举个最直观的例子一个 3×3 卷积输入 256 通道、输出 256 通道参数量是 256×256×9 ≈ 59 万而输入 64、输出 64 的同尺寸卷积只有 3.7 万差了 16 倍。Unet 越往深通道越多参数自然就堆起来了。显存占用还要看特征图本身。训练时前向的中间激活都要留着做反向传播一个 batch 里每层特征图的大小是批大小 × 通道数 × 宽 × 高。以输入 512×512、第一层 64 通道为例单张图的这一层激活就是 64×512×512 ≈ 1600 万个浮点数用 float32 存要 64MB批大小 4 就是 256MB还只是第一层。所以医学图像分割训练时batch size 常常只能开到 2 到 4不是不想开大是显存真放不下。优化方向具体做法代价减小输入尺寸从 512 降到 256 再上采样标签小目标细节损失降低基础通道64 改为 32 起步精度略降速度翻倍减少批大小从 8 降至 2收敛变慢BN 不稳定混合精度训练使用 fp16/bf16需注意数值溢出4.2 感受野是怎么一层层累加的感受野就是某个输出像素在原图上看到的范围。两个 3×3 卷积叠加的感受野是 5×5三个叠起来是 7×7规律是每多一层 3×3 就加 2。池化会让感受野翻倍因为尺寸缩小了后续卷积在原图上的等效覆盖范围就扩大了。感受野的意义在于它决定了网络能不能看到足够大的上下文。分割一个细胞如果感受野只覆盖细胞内部而看不到周围组织就很难判断边界。Unet 通过四层下采样把感受野堆到很大这也是它看起来不深、但效果不错的原因之一。不过感受野不是越大越好过大反而会把无关区域卷进来产生误判。4.3 为什么输入尺寸最好是 16 的倍数因为四次池化每次把尺寸整除 2累计要能被 16 整除否则某一级池化后尺寸会变成奇数下一级的拼接就会对不上。比如输入 500第三次池化后得到 62再池化得到 31上采样回来是 62而编码器对应层是 62恰好能对上但某些实现里会出现 31 除以 2 变成 15向下取整再上采样得到 30和编码器的 31 差了 1拼接直接报错。为省掉这类麻烦直接把输入 resize 到 16 的倍数最省心。如果你的任务需要保持长宽比可以先等比缩放再对短边做 padding 补齐到 16 的倍数预测完再把 padding 裁掉。这是我自己踩过几次坑之后固定下来的流程比事后改网络结构简单得多。5. 训练 Unet 时最容易翻车的几个地方5.1 数据量小到几百张时该怎么补医学分割数据集常常只有几十到几百张标注图直接拿去训 Unet 一定会过拟合训练集 dice 冲到 0.95验证集卡在 0.6 死活不涨。这时候数据增强不是可选项而是必需品。原始 Unet 论文特别强调了一种增强方式——弹性形变也就是在图像上施加随机的非线性扭曲。理由很直接生物组织本身就柔软、会形变这种增强和真实场景的分布最接近。除此之外随机旋转、水平垂直翻转、随机缩放、亮度和对比度抖动、随机裁剪都值得加。但要注意医学图像里的翻转并不总是安全的比如某些左右器官语义不同盲目翻转会让标签语义错乱。这一点在做肝脏、肺部分割时无所谓做左右手、单侧结构时就要小心。注意增强之后的标签必须和图像用完全相同的变换参数否则分割掩码和图片就对不上训出来的模型会学得莫名其妙。用同一套随机种子驱动图像和掩码的变换是最稳的做法。5.2 类别极端不平衡时损失函数该换哪个分割任务里前景像素常常只占全图 5% 甚至更少用普通交叉熵会有一个致命问题模型只要全部预测背景就能拿到很低的损失于是它学会了偷懒,永远不敢预测前景。表现就是训练 loss 一路下降但预测结果全是黑图。解决方案有两类。一类是加权交叉熵给前景像素更大的权重另一类是直接换用对类别比例不敏感的损失比如 Dice loss。Dice 系数衡量的是预测和真值两个集合的重合程度公式是 2×交集除以两边之和它不看背景有多少只看前景重叠得好不好所以对小目标是天然的照顾。实际用的时候我会用 Dice loss 打底再叠加一部分加权交叉熵或者 BCE让训练既有 Dice 那种对不平衡的鲁棒性又有交叉熵那种稳定的梯度。纯 Dice 训练有时会不太稳特别是训练初期预测全是空的时候分母接近零梯度容易抖。这个组合我自己在好几个分割任务上试下来都比单独用一个更顺。5.3 边界模糊与标签噪声分割任务里最难的永远是边界。原始 Unet 论文为此设计了一套边界加权方案给靠近细胞边界的像素更高的损失权重逼着网络把注意力放到难分的地方。它的权重公式由两部分组成一部分是平衡类别数量的固定权重另一部分随像素到边界距离指数衰减离边界越近权重越高。论文里那组参数是 w0 取 10、sigma 取 5。但加权是把双刃剑。如果标注本身有噪声比如两个标注者对同一条边界的判断差了好几个像素那么加权损失会放大这种噪声模型反而学偏。所以标注质量差的数据集上我更倾向于弱化边界权重或者先用一份质量较高的子集训练再用它去辅助修正其他样本。另外边界模糊还有一个常被忽略的来源——标签的下采样。如果你的标注是原图分辨率但训练时把图和标签一起缩放到 256×256边界细结构会在缩放中糊掉模型再厉害也学不回不存在的细节。能保持原分辨率训练就别缩。6. 从原始 Unet 到现在的各种变体改进到底改在哪6.1 换骨干网络把编码器换成更强的主干最主流的一类改进就是保持 U 形结构不动把编码器换成 ResNet、EfficientNet、VGG 这类在 ImageNet 上预训练过的主干。这么做的好处很实在——预训练权重给了网络一个很好的初始化小数据集上收敛快、效果稳。原始 Unet 的编码器是随机初始化的全靠数据自己学数据一少就学不好。换了骨干之后要注意通道数对齐的问题。ResNet 的各级输出通道和原始 Unet 的 64、128、256、512 不一定一致拼接前通常要加一个 1×1 卷积把通道调成一致或者干脆把拼接换成在通道维度上更灵活的融合方式。另外预训练主干的下采样倍数可能不同需要重新核对每一级的尺寸别想当然地认为和原始 Unet 一样是 16 倍。6.2 注意力机制与 Transformer 的引入第二类改进是往跳跃连接或解码器里塞注意力模块。典型的做法是在编码器特征传给解码器之前先给它算一个注意力权重抑制和当前类别无关的区域、突出相关区域。这在多类别分割上收益明显因为浅层特征里混着很多其他器官的边缘直接拼过去会干扰当前类的判断。近几年更大的改动是把 Transformer 引入分割。它在瓶颈层这种低分辨率、高语义的位置替换掉卷积用自注意力去建模长距离依赖。好处是全局感受野不再依赖堆叠卷积坏处是参数量和显存需求陡增而且对数据量更敏感。所以实践中常见的是混合结构——浅层还是卷积、只在最深的一两级用 Transformer既拿到全局建模又不至于让数据量和显存双双爆表。6.3 轻量化、多尺度输出与自动配置第三类改进方向更偏工程落地。移动端部署要用轻量主干把标准卷积换成深度可分离卷积参数量能降一个数量级代价是精度小幅下滑通常靠更多的数据增强和蒸馏补回来。多尺度输出则是让网络在多个分辨率上都产生预测再把它们加权融合对小目标和大小差异大的场景比较友好。还有一个不得不提的方向是自动配置框架它不发明新结构而是把预处理、网络结构、训练参数、后处理整套流程根据数据集特性自动推导出来在很多分割基准上长期压着各种手工设计的模型。它给人的启示其实很重要在分割任务里把数据分布摸清楚、把流程调对往往比换一个更花哨的模块更管用。我自己做项目时也是这个顺序——先把输入尺寸、归一化、增强、损失这些基础项对齐再考虑加注意力或者换主干顺序反了很容易白忙一场。模型训完只是开始推理阶段还有一件事值得留意滑动窗口和重叠拼接。整张大图直接塞进网络会因为显存和尺寸限制被压缩边界就会糊。用带重叠的滑窗逐块推理再拼接能明显改善大图上的连续性代价是推理时间变长。如果时间允许重叠比例给到窗口尺寸的 1/4 到 1/2边缘处理会干净很多。踩过几次边界错位的坑之后这一条我基本每次都会加上。