1. 为什么选择复现 UNet 而不是直接调库搞过语义分割的人都知道UNet 这个结构在医学影像、遥感、工业质检这些领域几乎是绕不开的。我第一次接触它是在一个缺陷检测的小项目里当时想偷懒直接调segmentation_models_pytorch里的现成模型结果发现预训练权重跟我的数据域差太远效果还不如自己从头搭一个。后来索性花了两天时间把原始论文的架构完整复现了一遍反而把很多细节吃透了。这篇东西就是那次复现的完整记录。我会从架构设计、数据管道、损失函数、训练循环到推理部署把每个环节的坑和取舍都讲清楚。适合已经会写 PyTorch 基础代码、想深入理解 UNet 内部机制的人也适合那些调库调不出效果、准备自己动手的从业者。读完你至少能拿到一份可以直接跑通的代码骨架以及一份别人不会告诉你的调参经验。先说结论UNet 的核心价值不在于它有多深而在于它的编码器-解码器对称结构加上跳跃连接。编码器负责逐层提取语义特征解码器负责逐层恢复空间分辨率而跳跃连接把编码器的高分辨率浅层特征直接拼到解码器的对应层级上。这个设计解决了一个根本矛盾——深层特征语义强但位置信息模糊浅层特征位置准但语义弱。拼接之后网络同时拥有“是什么”和“在哪里”两种信息。很多人复现的时候只关注通道数变化忽略了跳跃连接的具体拼接方式、上采样的实现细节、以及 padding 策略对输出尺寸的影响。这些细节才是决定你能不能复现出论文效果的关键。2. 整体架构拆解与设计取舍2.1 原始论文架构的逐层分析原始 UNet 论文里的结构其实很规整。编码器有四层下采样每层包含两个 3x3 卷积加 ReLU然后一个 2x2 最大池化。通道数从 64 开始每下采样一次翻倍64、128、256、512。最底层是 1024 通道的瓶颈层同样两个 3x3 卷积。解码器对称地有四层上采样。每层先做一个 2x2 转置卷积把特征图尺寸翻倍、通道数减半然后把编码器对应层的特征图裁剪到相同尺寸后拼接再接两个 3x3 卷积。最后用一个 1x1 卷积把 64 通道映射到类别数。这里有个容易被忽略的点原始论文用的是valid padding也就是卷积不补零。这意味着每次 3x3 卷积都会让特征图缩小 2 个像素。所以编码器每层的特征图尺寸是 568、280、136、64、28解码器需要从 28 恢复到 388最后再裁剪到 388x388 的输出。这个裁剪操作在论文里写得很清楚但很多人用 same padding 复现时就把它省掉了。我个人的建议是如果你用 same padding就不需要裁剪直接拼接即可。这样代码更简洁输出尺寸也跟输入一致实际项目中更实用。但你要知道原始论文不是这么做的面试被问到的时候别答错。2.2 编码器主干的选择从零搭建还是用预训练这是复现时第一个要做的决策。原始论文是从零训练的因为当时的医学影像数据集规模不大而且灰度图跟 ImageNet 的 RGB 分布差异明显。但现在情况不一样了如果你做的是自然图像分割用 ResNet 或 EfficientNet 做编码器主干、加载 ImageNet 预训练权重收敛速度和最终精度都会好很多。我的经验是分两种情况数据量小于 2000 张强烈建议用预训练编码器。从零训练几乎必然过拟合除非你做大量数据增强。数据量大于 10000 张且域差异大比如医学 CT、遥感多光谱从零训练可能更好因为预训练权重的纹理偏好可能反而是种干扰。如果你决定从零搭建就按原始论文的通道数来。如果要用预训练主干推荐resnet34或efficientnet-b0前者结构简单容易改后者参数量小推理快。下面是我常用的一个从零搭建的编码器块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, 3, padding1, biasFalse), nn.BatchNorm2d(mid_ch), nn.ReLU(inplaceTrue), nn.Conv2d(mid_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)注意我加了 BatchNorm。原始论文用的是无 BN 的版本但实测下来加 BN 训练更稳、收敛更快尤其是 batch size 较小的时候。如果你要严格复现论文把 BN 去掉即可但精度可能会掉一两个点。2.3 上采样的三种实现方式对比解码器的上采样有三种常见做法各有优劣方式实现优点缺点转置卷积nn.ConvTranspose2d可学习参数表达能力强容易产生棋盘格伪影双线性插值F.interpolate无参数平滑表达能力弱最近邻插值F.interpolate简单快速块状效应明显原始论文用的是转置卷积。但我实测下来双线性插值加一个 3x3 卷积的组合往往效果更稳定而且参数量更少。棋盘格伪影在转置卷积里是个真实存在的问题尤其是 stride 和 kernel size 不匹配的时候。如果你非要用转置卷积确保kernel_size能被stride整除比如 kernel2、stride2或者 kernel4、stride2。我现在的默认选择是双线性插值代码大概长这样class UpBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv DoubleConv(in_ch skip_ch, out_ch) def forward(self, x, skip): x self.up(x) x torch.cat([skip, x], dim1) return self.conv(x)align_cornersTrue这个参数值得说一下。当设为 True 时插值会把输入特征图的角点对齐到输出角点尺寸关系是out (in-1)*scale 1。设为 False 时是out in*scale。对于 2 倍上采样如果输入是 28True 得到 55False 得到 56。这会影响跟跳跃连接拼接时的尺寸匹配所以要么统一用 False要么在拼接前做一次裁剪或 padding。我踩过这个坑当时编码器输出 56、解码器上采样出 55直接报错。3. 数据管道与损失函数的实战细节3.1 数据增强策略别照搬分类任务的那一套语义分割的数据增强跟图像分类有本质区别——你不仅要变换图像还要同步变换掩码。而且有些增强对分类有效对分割反而有害。我常用的增强组合是这样的随机水平翻转安全几乎总是有益。随机旋转 90 度安全适合医学和遥感图像。随机缩放裁剪有效但要注意裁剪后可能丢失小目标。颜色抖动只对图像做掩码不动。强度要小否则可能改变类别语义。弹性形变医学图像常用但计算开销大慎用。有一个坑我必须提醒不要用 RandomResizedCrop 这种会改变目标形状的增强除非你的任务对形状不敏感。分割任务里目标的边界信息很重要随意缩放裁剪会破坏空间一致性。另外掩码的插值方式必须用最近邻不能用双线性。因为标签是整数类别双线性插值会产生 0.5 这种无意义的类别值。这个错误我见过不止一个新手犯。import albumentations as A train_transform A.Compose([ A.HorizontalFlip(p0.5), A.RandomRotate90(p0.5), A.ShiftScaleRotate(shift_limit0.0625, scale_limit0.1, rotate_limit15, p0.5), A.ColorJitter(brightness0.2, contrast0.2, p0.3), ], additional_targets{mask: mask})用 albumentations 的好处是它自动处理图像和掩码的同步变换省得你自己写。如果你不想引入额外依赖用 torchvision 的transforms配合手动同步也行但代码会啰嗦不少。3.2 损失函数选择Dice、BCE 还是组合UNet 原始论文用的是逐像素的交叉熵。但在类别极度不平衡的分割任务里比如缺陷检测缺陷像素可能只占 1%交叉熵会被背景类主导模型倾向于全预测背景。这时候 Dice Loss 就派上用场了。Dice 系数衡量的是预测和真实掩码的重叠程度对类别不平衡不敏感。但 Dice Loss 也有缺点训练初期梯度不稳定容易震荡。我的经验是组合使用BCE Dice或者CE Dice。BCE 提供稳定的逐像素梯度Dice 负责拉高整体重叠度。权重一般设成 1:1 或者 1:2具体看你的类别不平衡程度。class DiceLoss(nn.Module): def __init__(self, smooth1e-6): super().__init__() self.smooth smooth def forward(self, logits, targets): probs torch.sigmoid(logits) probs probs.view(-1) targets targets.view(-1) intersection (probs * targets).sum() dice (2. * intersection self.smooth) / (probs.sum() targets.sum() self.smooth) return 1 - dice class ComboLoss(nn.Module): def __init__(self, bce_weight1.0, dice_weight1.0): super().__init__() self.bce nn.BCEWithLogitsLoss() self.dice DiceLoss() self.bce_weight bce_weight self.dice_weight dice_weight def forward(self, logits, targets): return self.bce_weight * self.bce(logits, targets) \ self.dice_weight * self.dice(logits, targets)多分类的情况把 BCE 换成 CrossEntropyDice 改成按类别计算再平均。注意 softmax 和 sigmoid 的区别别搞混了。3.3 学习率调度与优化器配置优化器我基本只用 Adam 或 AdamW。SGD 虽然理论上泛化更好但收敛太慢调参成本高。AdamW 比 Adam 多了权重衰减的解耦对大模型更友好UNet 这种规模用哪个都行。学习率初始值设 1e-3 或 3e-4。如果用了预训练编码器编码器部分的学习率要设小一点比如 1e-4解码器用 1e-3。这叫判别式微调能防止预训练权重被快速破坏。调度器推荐CosineAnnealingLR或ReduceLROnPlateau。前者平滑下降后者根据验证集指标动态调整。我一般用余弦退火配合 warmup前 5 个 epoch 从 1e-6 线性升到初始学习率然后再余弦下降。这样训练初期不会因为随机初始化的解码器产生大梯度而震荡。from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) warmup LinearLR(optimizer, start_factor1e-3, total_iters5) cosine CosineAnnealingLR(optimizer, T_max95, eta_min1e-6) scheduler SequentialLR(optimizer, schedulers[warmup, cosine], milestones[5])4. 完整训练流程与关键环节实现4.1 模型组装与尺寸校验把编码器、瓶颈层、解码器拼起来的时候最容易出错的就是尺寸匹配。我建议在forward里加断言训练前先用一个假输入跑一遍确认每层输出尺寸符合预期。class UNet(nn.Module): def __init__(self, in_ch3, num_classes1, features[64, 128, 256, 512]): super().__init__() self.downs nn.ModuleList() self.ups nn.ModuleList() self.pool nn.MaxPool2d(2, 2) # 编码器 for f in features: self.downs.append(DoubleConv(in_ch, f)) in_ch f # 瓶颈层 self.bottleneck DoubleConv(features[-1], features[-1]*2) # 解码器 for f in reversed(features): self.ups.append(UpBlock(f*2, f, f)) self.final_conv nn.Conv2d(features[0], num_classes, 1) def forward(self, x): skips [] for down in self.downs: x down(x) skips.append(x) x self.pool(x) x self.bottleneck(x) skips skips[::-1] for up, skip in zip(self.ups, skips): x up(x, skip) return self.final_conv(x)跑一遍model(torch.randn(2, 3, 256, 256))确认输出是(2, num_classes, 256, 256)。如果尺寸对不上八成是上采样和池化的尺寸计算有偏差。256 这种 2 的幂次尺寸最安全因为每次池化都整除。如果你用 250 这种非 2 幂次编码器最后会变成 15、7、3 这种奇数上采样回来就对不上了。建议输入尺寸统一 resize 到 32 的倍数比如 256、320、512。4.2 训练循环的骨架与混合精度训练循环本身不复杂但有几个细节决定成败。第一是验证集指标要跟损失分开看损失下降不代表分割质量提升一定要盯着 Dice 或 IoU。第二是保存最佳模型而不是最后一个 epoch 的模型过拟合往往发生在后期。混合精度训练AMP能省显存、加速训练现在 PyTorch 支持得很好几乎无脑开就行。但要注意 Dice Loss 在 fp16 下可能数值不稳定建议在损失计算时强制转回 fp32。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for epoch in range(num_epochs): model.train() for imgs, masks in train_loader: imgs, masks imgs.cuda(), masks.cuda() optimizer.zero_grad() with autocast(): logits model(imgs) loss criterion(logits.float(), masks.float()) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step() # 验证 model.eval() dice_scores [] with torch.no_grad(): for imgs, masks in val_loader: imgs imgs.cuda() logits model(imgs) preds (torch.sigmoid(logits) 0.5).float() dice_scores.append(compute_dice(preds, masks.cuda())) print(fEpoch {epoch}, Dice: {sum(dice_scores)/len(dice_scores):.4f})autocast上下文里的前向计算用 fp16但损失计算我手动转了 fp32。这个细节很多人忽略结果训练到一半 loss 变成 nan还找不到原因。4.3 推理与后处理阈值选择比模型本身还重要训练完之后推理阶段最关键的参数是二值化阈值。默认 0.5 不一定最优。如果模型偏向保守预测正类少可以把阈值降到 0.3 或 0.4如果偏向激进就升到 0.6。我通常会在验证集上扫一遍阈值从 0.1 到 0.9步长 0.05找 Dice 最高的那个。这个操作花不了几分钟但往往能带来一两个点的提升。后处理还包括连通域过滤去掉面积过小的预测区域它们多半是噪声。用cv2.connectedComponents或者scipy.ndimage.label都能做。面积阈值根据你的目标尺寸来定一般设成最小真实目标面积的 1/4 左右。import numpy as np from scipy import ndimage def postprocess(pred_mask, min_area50): labeled, num ndimage.label(pred_mask) for i in range(1, num1): if (labeled i).sum() min_area: pred_mask[labeled i] 0 return pred_mask还有一个技巧是测试时增强TTA对输入做水平翻转、多尺度缩放分别推理后把结果平均。能稳定提升一个点左右代价是推理时间翻几倍。对精度要求高的场景值得做。5. 常见问题与排查技巧实录5.1 训练不收敛或 loss 震荡这是最常见的问题。排查顺序我一般是这样的检查数据标签掩码是不是全黑或者全白类别值是不是从 0 开始连续我遇到过标签是 0 和 255 的情况模型直接学废了。检查学习率太大导致震荡太小导致不下降。先用 1e-3 试不行降到 1e-4。检查损失函数BCE 的输入必须是 logits不能再过 sigmoid。过了 sigmoid 再喂给 BCEWithLogitsLoss 会导致梯度消失。检查 BatchNormbatch size 小于 8 的时候 BN 统计量不准建议换成 GroupNorm 或者 InstanceNorm。下面这张表是我整理的高频问题速查现象可能原因解决方法loss 不下降学习率过小、标签错误调大 lr、检查标签loss 震荡学习率过大、batch 太小调小 lr、增大 batch验证集 Dice 远低于训练过拟合加数据增强、加 dropout、早停预测全为背景类别不平衡用 Dice Loss、调高正类权重预测边界粗糙上采样方式问题换双线性插值、加 CRF 后处理显存溢出batch 太大、模型太深减小 batch、用 AMP、梯度累积5.2 显存不够用的几种解法UNet 在 512x512 输入下batch size 开到 8 就可能爆显存。除了换小模型还有几个实用技巧梯度累积batch size 设 2累积 4 次再更新等效 batch size 8。混合精度直接省一半显存。检查点机制torch.utils.checkpoint用计算换显存适合超深网络。减小输入尺寸从 512 降到 384显存占用降 44%精度可能只掉一个点。梯度累积的代码很简单accum_steps 4 for i, (imgs, masks) in enumerate(train_loader): loss criterion(model(imgs), masks) / accum_steps loss.backward() if (i1) % accum_steps 0: optimizer.step() optimizer.zero_grad()注意 loss 要除以累积步数否则梯度会放大。5.3 从单类分割扩展到多类的注意事项单类分割输出通道是 1用 sigmoid。多类分割输出通道是类别数用 softmax。这个切换有几个坑损失函数从 BCE 换成 CrossEntropyDice 要按类别算再平均。掩码的类别值必须是 0 到 num_classes-1 的整数不能有 255 这种忽略值除非你在损失里显式处理。推理时用argmax而不是阈值二值化。如果类别不平衡严重考虑用带权重的 CrossEntropy 或者 Focal Loss。我做过一个五类分割的项目背景占 90%四个目标类各占 2.5% 左右。直接用 CE 的话模型只学背景就能到 90% 准确率但 Dice 惨不忍睹。后来换成 CE 按类别 Dice并且给稀有类加权才把效果拉上来。6. 几个让我印象深刻的踩坑经历第一次复现的时候我完全照搬论文的 valid padding结果输入 572x572 输出 388x388跟我的标签尺寸对不上。当时不知道要裁剪硬是把标签也 resize 到 388结果边界信息全丢了。后来才明白要么用 same padding 保持尺寸要么在拼接时把编码器特征裁剪到解码器尺寸。这个教训让我养成了每次拼接前打印尺寸的习惯。还有一次训练集 Dice 到 0.95验证集只有 0.6。我以为是过拟合加了 dropout 和权重衰减都没用。后来发现是数据泄露——训练集和验证集有重叠样本。重新划分数据后验证集 Dice 直接到 0.85。所以划分数据集之前一定要去重尤其是从视频里抽帧得到的数据。最后一个坑是关于 BatchNorm 的。有次我用 batch size 2 训练BN 的 running mean 和 variance 估计得极差推理时效果比训练时差一大截。换成 GroupNorm 后问题消失。所以小 batch 场景下别用 BN用 GN 或 IN。这些经验在官方文档里都不会写但实际项目中每一个都能让你卡半天。复现 UNet 不难难的是把这些细节都处理对。希望这份记录能帮你少走点弯路。