问十个做过图像分割的人至少八个第一次正经接触的模型是UNet。这话不夸张当年我在一个医疗影像方向的项目里需要把显微镜图像里的目标区域逐像素圈出来翻了一圈资料之后最终落地的baseline还是UNet。倒不是说它最先进而是这个网络结构足够清晰、改动空间大、在小数据集上也能训练得动。Pytorch实现UNet的开源代码一抓一大把但我见过不少人照着抄完却跑不出预期效果问题大多不在网络主体而在跳连对齐、损失函数、数据组织这类细节上。这篇文章会从零搭一个可运行的UNet把每个模块的代码、背后的设计逻辑以及我在实际项目中踩过的坑一次讲清楚。1. 从分割任务说起UNet为什么长这样1.1 语义分割到底在解决什么问题先弄清楚任务本身。图像分类的输出是一个标签告诉你这张图里有什么目标检测的输出是若干边界框告诉你东西在哪一块区域而语义分割要的是更细的结果——给每一个像素都贴上类别标签。输出是一张和输入尺寸相同的掩码图每个像素点的取值代表它属于哪一类。这个任务有几个天然难点。一是边界精度模型不光要知道物体大致在哪个区域还要把边缘轮廓勾勒清楚二是小目标比如医学图像里早期的病灶可能只有几十个像素很容易被下采样吞掉三是类别不平衡很多分割任务里背景像素占了九成以上模型随便输出全背景都能拿到很高的准确率但这显然不是我们要的结果。UNet之所以成为分割领域的常青树恰恰是在这几个问题上给出了一个工程上非常有效的方案。1.2 编码-解码结构与跳连的直觉解释UNet的整体结构可以拆成三条线编码器、解码器、跳连。编码器部分做的事情是不断下采样。通过卷积提取特征再用最大池化把分辨率降下来一路把通道数从64翻到128、256、512甚至1024。这个过程很像读一篇长文章时先看大标题、再抓段落主旨——高层特征回答的是这里大概是什么感受野越来越大空间位置信息却越来越模糊。解码器部分做的事情正好反过来。通过上采样把分辨率逐步恢复同时通道数逐步减半最终回到原图尺寸。解码器负责回答这个区域具体长什么样需要把粗糙的高层语义信息还原成精细的空间结构。问题来了如果只靠编码器到最后的高层特征去还原细节早就丢光了。UNet的聪明之处在于跳连——把编码器每一层下采样之前得到的特征图直接拼接到解码器对应层上。这样解码器既能看到全局语义又能拿到浅层细节相当于画地图时先看整体布局再对着卫星照片把每条小巷子补齐。道理讲完下面动手写代码。2. 动手实现从DoubleConv到完整网络2.1 DoubleConvUNet最基础的积木UNet里出现频率最高的结构是连续两次3x3卷积 激活函数论文里叫double convolution。我这里加上了BatchNorm说句实话原论文里是没有BN的但现代数据集和训练配置下加上BN通常更稳定收敛也更快。先看代码import torch import torch.nn as nn class DoubleConv(nn.Module): 两次 3x3 卷积 BN ReLUUNet 的基本组成单元 def __init__(self, in_channels, out_channels): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)两个3x3卷积叠加等价于一个5x5的感受野但参数量更少2×918个参数权重比25个小不少而且中间多了一次非线性激活表达能力反而更强。这里把bias设为False是因为BatchNorm自带可学习的偏移项如果卷积层再加bias信息是冗余的还容易在初始化阶段出问题。2.2 Down与Up编码器与解码器的骨架编码器的每一步是最大池化 DoubleConv解码器的每一步是上采样 拼接 DoubleConv。上采样我用的是转置卷积ConvTranspose2dkernel_size和stride都设为2可以把特征图尺寸放大一倍、通道数减半。这也是论文里up-conv的标准做法。class Down(nn.Module): 下采样2x2 最大池化 DoubleConv def __init__(self, in_channels, out_channels): super().__init__() self.maxpool nn.MaxPool2d(kernel_size2, stride2) self.conv DoubleConv(in_channels, out_channels) def forward(self, x): return self.conv(self.maxpool(x)) class Up(nn.Module): 上采样转置卷积放大 跳连拼接 DoubleConv def __init__(self, in_channels, out_channels): super().__init__() self.up nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size2, stride2) self.conv DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 self.up(x1) # 处理输入尺寸不能被2整除时产生的尺寸错位 diffY x2.size()[2] - x1.size()[2] diffX x2.size()[3] - x1.size()[3] x1 nn.functional.pad( x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2] ) x torch.cat([x2, x1], dim1) return self.conv(x)Up模块里有几个坑要解释清楚。第一个是参数命名forward接收两个输入x1是解码器传来的深层特征x2是编码器对应层的跳连特征。转置卷积先把x1放大然后和x2在通道维度上拼接——注意是拼接不是相加。第二个是尺寸对齐这个在下一节详细说。2.3 组装完整网络并验证每一层的尺寸有了积木组装UNet就很简单了。以经典的四次下采样为例输入输出尺寸全程通过padding保持对齐不需要像原论文那样用valid卷积把572x572变成388x388写起来也直观得多class UNet(nn.Module): def __init__(self, in_channels3, num_classes1, base_channels64): super().__init__() self.in_channels in_channels self.num_classes num_classes self.inc DoubleConv(in_channels, base_channels) # 64 self.down1 Down(base_channels, base_channels * 2) # 128 self.down2 Down(base_channels * 2, base_channels * 4) # 256 self.down3 Down(base_channels * 4, base_channels * 8) # 512 self.down4 Down(base_channels * 8, base_channels * 16) # 1024 self.up1 Up(base_channels * 16, base_channels * 8) # 512 self.up2 Up(base_channels * 8, base_channels * 4) # 256 self.up3 Up(base_channels * 4, base_channels * 2) # 128 self.up4 Up(base_channels * 2, base_channels) # 64 self.outc nn.Conv2d(base_channels, num_classes, kernel_size1) def forward(self, x): x1 self.inc(x) # 1 倍分辨率, 64 通道 x2 self.down1(x1) # 1/2 分辨率, 128 通道 x3 self.down2(x2) # 1/4 分辨率, 256 通道 x4 self.down3(x3) # 1/8 分辨率, 512 通道 x5 self.down4(x4) # 1/16 分辨率, 1024 通道 x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) logits self.outc(x) return logits写完之后第一件事就是用假数据跑一遍前向确认每个张量的尺寸符合预期。这是我最想强调的习惯设计网络时先做形状检查别等训练到一半才被报错打断。model UNet(in_channels3, num_classes1, base_channels64) fake torch.randn(2, 3, 256, 256) out model(fake) print(out.shape) # 期望输出 torch.Size([2, 1, 256, 256])参数规模方面base_channels64的经典UNet大约3100万参数。这个数字意味着用Adam优化器训练时优化器状态本身的内存几乎是参数的三倍再加上中间特征图8G显存跑256x256的batch size可能只敢设到48。显存吃紧的同学可以把base_channels降到32参数量直接缩到四分之一效果并不会差太多这是个非常实用的调参手段。3. 跳连与尺寸对齐最容易被忽略的细节3.1 为什么跳连选择拼接而不是相加这是个值得停下来想清楚的问题。ResNet里的残差连接用的是相加UNet的跳连用的是拼接同样是融合浅层和深层信息为什么做法不一样关键在通道语义上。UNet的编码器每一层都在改变通道数64、128、256、512、1024。如果采用相加要求两个特征图的通道数必须一致这就等于强行把编码器某一层的信息压缩成和解码器相同的通道规模再合并浅层特征里那些空间细节很可能在这一步就被折损了。拼接则完全保留了双方的信息——解码器拿到拼接结果后由DoubleConv里的卷积层自己去学习该从浅层特征里挑什么、从深层特征里取什么。换句话说相加是让网络在已有信息上做修正拼接是让网络在更完整的信息上做筛选后者对定位细节双重需求的分割任务更友好。3.2 输入尺寸不能整除16时的处理方案UNet下采样到最深处时分辨率变成输入的1/16这意味着输入尺寸必须是16的倍数整个网络的编解码尺寸才能严丝合缝地对上。实际项目里哪有那么巧原始图片可能是512x512、1000x1000更常见的是像572x572这类随意尺寸池化之后必然出现奇数尺寸转置卷积放大后和跳连特征就会差一个像素。我的处理方法是上面Up模块里的写法先计算两个特征图的高度差和宽度差把放大后的特征 pad 到和跳连特征一致再拼接。这样最终输出尺寸能严格等于输入尺寸。另一个常见做法是直接裁剪跳连特征x2 x2[:, :, :x1.size(2), :x1.size(3)]这其实是原论文里的思路——跳连的特征更大那就丢掉边界的一些像素。两种方案都能跑我的习惯是pad因为裁剪会丢掉一个像素的信息而pad只是补零对结果影响更可控。不过记住一条铁律训练和推理时输入尺寸的处理方式必须完全一致否则模型在训练时学到的空间对应关系推理时会对不上。3.3 上采样方式的选择转置卷积还是插值前面代码里用的是转置卷积这是UNet原版的方案。它的好处是上采样的细节是学出来的表达能力更强。缺点也常被提到kernel_size如果和stride配合不好输出会产生棋盘格状的伪影因为转置卷积在重叠区域会出现不均匀的加权。实际操作中很多人会把上采样改成双线性插值/最近邻插值 3x3卷积代码上就是把Up里的self.up换成self.up nn.Sequential( nn.Upsample(scale_factor2, modebilinear, align_cornersTrue), nn.Conv2d(in_channels, in_channels // 2, kernel_size3, padding1, biasFalse) )两种方案我都在项目里跑过。转置卷积在医学影像这类边界相对柔和的任务上表现不错插值卷积几乎没有棋盘格问题参数也少收敛更稳。如果你刚上手我建议先用插值卷积把整套流程跑通之后再换成转置卷积对比一下哪个效果好留哪个。别纠结于哪个更正统分割任务里以结果论英雄。4. 训练配置损失函数、数据增强与评测指标4.1 多损失组合BCE Dice的实践网络输出的是未经激活的logits二分类分割任务里最常见的损失是BCEWithLogitsLoss它内部把Sigmoid和交叉熵合并计算数值上比分开调用更稳定。但只用BCE在背景占比极高时模型很容易偏向全预测为背景因为这样损失已经很小了。所以我通常把Dice Loss加进去一起优化import torch.nn.functional as F def bce_dice_loss(logits, target, smooth1e-6): bce F.binary_cross_entropy_with_logits(logits, target) prob torch.sigmoid(logits) # target 和 prob 都视为 [B, N] 的展平向量 intersection (prob * target).sum() dice 1 - (2.0 * intersection smooth) / (prob.sum() target.sum() smooth) return bce dice为什么这么组合BCE是逐像素独立的对每个点的错误都敏感梯度信号密集Dice看的是整体重叠程度直接优化预测区域和真实区域的交并比对类别不平衡不那么敏感。两者互补后模型既不会忽略小的目标也不会因为类别失衡而停滞。多分类场景就把BCE换成CrossEntropyLossDice部分区分每个类别单独计算再取平均。4.2 训练数据增强的三个注意点分割任务尤其是医疗方向样本量往往几百张甚至几十张不做增强根本没法训练。最常用的是随机水平翻转、垂直翻转、随机旋转、随机裁剪缩放。但这里有个非常容易踩的坑图像和掩码必须同步变换。很多人用Pytorch自带的torchvision.transforms写了一堆随机翻转、旋转的变换结果发现只有图像变了、掩码没跟着变训练出来的模型自然是错误的。正确做法是固定随机种子再分别对image和mask做同样的操作或者直接用支持同时变换图像和掩码的增强库这类库内部会保证两者使用同一组变换参数。第二个注意点是弹性形变医学图像里组织本身有弹性适度的弹性形变增强能明显提升泛化能力。第三个注意点是不要把增强做得太狠比如90度、180度大幅旋转在细胞分割里可能合适在道路分割里就完全违背物理常识。增强策略要贴合数据本身的分布。4.3 评测指标为什么要用Dice和IoU而不是像素准确率训练时看损失评估时看指标两个概念别混。分割领域的两个标准指标是Dice系数和IoU。对二分类来说IoU TP / (TP FP FN) Dice 2*TP / (2*TP FP FN)两者高度相关Dice数值通常会比IoU高一些医学文献里更爱用Dice遥感、自动驾驶场景更爱用IoU。它们的共同好处是天然缓解了类别不平衡的影响——背景预测得再好如果前景区域对不上分数照样低。我见过一个例子背景占95%的数据集上像素准确率能到97%听起来很吓人但模型其实什么都没学会掩码画出来是一团糊。换成Dice一算可能只有0.2。所以在项目一开始就要约定好评测指标不要被准确率这种假繁荣误导。多分类任务里通常会报告每个类别的IoU再取平均mIoU小类别单独看指标不能只看平均值。5. 从跑通到可用调试中的坑与对策5.1 全背景预测问题模型一上来就摆烂我最早跑UNet时遇到的最烦人的问题是训练没几个epoch验证集上预测的掩码全是背景一个前景像素都没有。损失还在下降因为背景占了绝大多数模型发现什么都不预测就是当前损失最低的路线。逐个排查最常见的三个原因。第一是前景占比太小BCE里加上类别权重或者像我前面那样叠加Dice损失问题通常立刻缓解。第二是学习率设置得太激进模型在前几个batch里就冲进了全背景这个局部最优把初始学习率从1e-3调到1e-4基本能避免。第三是掩码预处理出错比如标签值不是0和1而是0和255或者mask的类型不是float32导致BCE计算时数值不对仔细检查target.unique()的结果这一步能省下好几个小时。5.2 BatchNorm在小batch下的表现UNet这种深层卷积网络BatchNorm在小batch下很不稳定。医学图像往往很大受显存限制batch只能设到1或2。这时候BN统计量在训练过程中剧烈抖动验证时经常出现预测结果忽好忽坏。三种解决办法一是尽量加大batch size哪怕把输入裁剪成小块二是换成GroupNorm或InstanceNorm它们不依赖batch维度batch1也能稳定训练三是把BN层换成nn.SyncBatchNorm配合多卡训练但单卡场景意义不大。我的经验是如果batch size长期小于4直接上GroupNorm比较省心。当然在显存允许的前提下优先保证batch size达到8以上BN的收益还是明显的。5.3 推理阶段的大图切片训练时用的输入是256x256的裁剪块但实际推理时面对的往往是整张大图直接resize会损失小目标的细节而且显存也放不下。业界标准的办法是滑窗切片推理把大图按固定stride切成若干patch每个patch单独过网络再把预测结果拼回完整尺寸。切片之间要有重叠因为patch边缘的预测质量通常比中心差重叠区域用加权平均或者只取每个patch中心区域的结果能显著减少拼接缝。如果追求更高精度还可以做测试时增强对输入做翻转、旋转后分别预测再把结果归一化平均。这个方法不用改模型就能稳定提升几个点的Dice缺点是推理时间成倍增加。实际项目中先看需求实时性要求高的场景就不要用了。5.4 我的几个实用习惯最后分享几个在实践中沉淀下来的习惯。第一训练前先用一小批数据跑一个batch的过拟合测试如果模型能在一个batch上把损失压到很低说明网络本身没问题后续训练失败的锅不在结构上。第二每个epoch结束都保存模型权重保留验证集Dice最好的那一份而不是只看最后一个epoch因为靠近收敛时指标容易反复震荡。第三记录一份固定的超参数基线256x256输入、base_channels64、Adam初始学习率1e-4、batch size 8、BCEDice组合损失在这个基线上做单变量调整出了问题永远知道该往哪里查。UNet的结构再过多少年都不会过时它是理解分割任务的最佳起点。把这套Pytorch实现吃透之后你再去接触DeepLab、SegNet或者Transformer系的分割模型会发现很多设计思路是相通的。项目里遇到分割需求我的第一反应依然是先跑一个UNet baseline——它简单、稳定、可解释而baseline的意义就是让你知道后面所有花哨的改进到底有没有真的在起作用。