1. U-Net 为什么值得反复拆解第一次接触 U-Net 是在做医学影像分割的项目里当时手头只有几百张标注图像试过几个常规的编解码结构效果都不太理想。后来换成 U-Net在同样数据量下 Dice 系数直接上了一个台阶。从那时起我就意识到这个 2015 年提出的结构之所以到现在还被大量引用和改造不是因为它复杂恰恰是因为它把“少样本、高精度、强定位”这三件事用最朴素的方式捏在了一起。U-Net 本质上是一个编码器-解码器结构编码器负责逐层下采样、提取从纹理到语义的抽象特征解码器负责逐层上采样、把特征图恢复到原图分辨率最终输出逐像素的分类结果。它最标志性的设计是跳跃连接把编码器每一层的特征图直接拼接到解码器对应层让浅层的高分辨率细节不至于在下采样过程中被彻底丢掉。这个思路听起来简单但在分割任务里极其关键——没有它边界会糊成一团有了它即便训练样本很少网络也能靠这些“抄近路”的特征把轮廓抠得比较准。这篇文章面向的是想真正把 U-Net 跑起来的人不管你是刚学完卷积神经网络、想找一个完整项目练手的学生还是需要在医学影像、遥感、工业质检等场景里落地分割模型的工程师我都会从架构设计讲到 PyTorch 代码实现把每一步的“为什么”说清楚。我不会只贴一段能跑的代码就完事而是把参数选择、数据组织、损失函数、训练技巧、常见报错都摊开讲让你看完能自己改、自己调、自己排查。2. 网络架构与核心原理拆解2.1 整体结构一个对称的“U”形是怎么来的U-Net 的名字来自它的形状。把网络画出来左边一路下采样右边一路上采样中间用跳跃连接横向连起来整体像一个字母 U。这个形状不是审美选择而是功能决定的。左侧编码器通常由若干个“卷积块 下采样”组成。每个卷积块一般是两次 3×3 卷积每次后面接 ReLU再跟一个 2×2 最大池化把空间尺寸减半、通道数翻倍。以输入 572×572 的灰度图为例经过几次下采样后特征图会变成 28×28 甚至更小但通道数从 64 涨到 512语义信息越来越浓缩。右侧解码器则反过来先上采样把尺寸放大一倍然后和左侧对应层的特征图拼接再做两次 3×3 卷积。最后一层用 1×1 卷积把通道数压到类别数输出分割图。这里有个细节很多人第一次看会忽略原始 U-Net 论文里输入和输出尺寸并不相等因为当年用的是 valid 卷积每卷一次尺寸就缩一点。现在我们在 PyTorch 里实现时通常用 padding1 的 same 卷积让输入输出尺寸保持一致这样数据组织会简单很多。这是基于当前主流实践的合理调整不是对原论文的否定。2.2 跳跃连接U-Net 的灵魂所在如果只能保留 U-Net 的一个设计那一定是跳跃连接。它的作用可以从两个角度理解。从信息流角度看编码器在下采样时虽然感受野变大、语义变强但每个像素对应的空间位置信息被不断稀释。到了最底层一个特征点可能对应原图上很大一块区域边界早就模糊了。解码器如果只靠这些高层特征往上恢复分割结果就会像用低分辨率图片放大一样边缘全是锯齿和毛刺。跳跃连接把编码器浅层的高分辨率特征直接送过来相当于给解码器提供了一份“边界参考图”。从梯度角度看跳跃连接还是一条短路径。反向传播时梯度可以从解码器直接流回编码器浅层不需要穿过整个深层网络。这对训练稳定性有帮助尤其是在数据量不大、网络又比较深的时候能缓解梯度消失的问题。实际操作中拼接方式一般是通道维度上的 concatenation而不是相加。相加会丢失一部分信息拼接则让网络自己决定怎么融合。代价是通道数会翻倍显存占用增加所以如果你的显卡比较紧张可以在拼接后加一个 1×1 卷积把通道压回去这是常见的轻量化改法。2.3 上采样的几种做法与选择逻辑解码器里的上采样有好几种实现方式选哪种会直接影响效果和速度。最常用的是转置卷积也叫反卷积。它通过学习参数来放大特征图理论上更灵活但容易产生棋盘格伪影尤其是在步长和卷积核尺寸不匹配的时候。另一种是双线性插值 普通卷积先插值放大再卷积融合效果通常更平滑参数也更少。还有一种更简单的最近邻插值速度最快但比较粗糙一般只在轻量场景用。我在实际项目里的经验是如果追求精度且显存充足用转置卷积如果追求稳定和速度用双线性插值加卷积如果只是做原型验证最近邻也能凑合。没有绝对优劣关键看你的任务对边界精度的要求有多高。2.4 损失函数分割任务不能只看准确率分割任务有个特点背景像素往往远多于前景像素。比如一张医学影像里病灶可能只占几个百分点。如果只用普通的交叉熵损失网络会倾向于把所有像素都预测成背景准确率看起来很高但实际什么都没分出来。所以 U-Net 类任务里常用的损失是Dice Loss或者交叉熵 Dice 的组合。Dice 系数衡量的是预测区域和真实区域的重叠程度对类别不平衡不敏感。它的定义是 2×交集 / (预测面积 真实面积)取值在 0 到 1 之间越接近 1 越好。作为损失使用时通常写成 1 - Dice。另一个常见选择是BCEWithLogitsLoss它把 Sigmoid 和交叉熵合在一起数值稳定性更好。如果是多类别分割就用 CrossEntropyLoss。我的习惯是二分类分割用 BCE Dice 各占一半权重多分类用 CrossEntropy Dice这样既能保证像素级分类准确又能优化区域重叠。3. PyTorch 代码实现与关键细节3.1 环境准备把 PyTorch 装明白在写模型之前先把环境弄干净。我见过太多人卡在安装这一步后面代码没问题却跑不起来很打击积极性。如果你用 Anaconda推荐单独建一个环境不要和 base 混在一起。命令大概是先创建一个 Python 3.9 或 3.10 的环境然后根据你的显卡情况选择安装命令。有 NVIDIA 显卡且驱动正常的话去 PyTorch 官网查对应 CUDA 版本的安装命令用 conda 或 pip 都行。没有显卡就用 CPU 版本学习阶段完全够用只是训练慢一些。注意不要盲目复制网上的安装命令CUDA 版本、Python 版本、操作系统三者要匹配。装完后用torch.cuda.is_available()验证一下返回 True 才说明 GPU 可用。验证环境是否正常可以跑一段极简代码创建一个随机张量做一次卷积看输出尺寸对不对。这一步能排除大部分环境问题。3.2 卷积块与下采样模块的写法U-Net 的编码器由重复的卷积块组成所以先把它封装成一个类。每个卷积块包含两次 3×3 卷积每次后面接 BatchNorm 和 ReLU。BatchNorm 能加速收敛、稳定训练ReLU 提供非线性。padding 设为 1保证卷积后尺寸不变。下采样用 2×2 最大池化步长为 2这样尺寸正好减半。为什么用最大池化而不是平均池化因为最大池化保留的是最显著的特征响应在分割任务里更有利于突出边界和纹理。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.block(x)这段代码看起来简单但有几个点值得说。inplaceTrue能省一点显存但如果你后面要复用 ReLU 前的值就别用。BatchNorm 在 batch size 很小时表现不稳定如果显存只够跑 batch size 为 1 或 2可以考虑换成 GroupNorm 或 InstanceNorm这是医学影像分割里常见的调整。3.3 编码器、解码器与跳跃连接的组装编码器就是四个“卷积块 池化”的堆叠通道数依次是 64、128、256、512。每次池化后特征图尺寸减半。最底层再过一个卷积块通道数到 1024这是整个网络的“语义瓶颈”。解码器每一步先上采样然后把编码器对应层的特征拼过来再过一个卷积块。拼接时要注意通道数对齐上采样后的通道数加上编码器送来的通道数正好是卷积块的输入通道数。class UNet(nn.Module): def __init__(self, in_ch1, out_ch2): super().__init__() self.enc1 DoubleConv(in_ch, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.enc4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(512, 1024) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.dec4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)这段代码可以直接跑但要注意输入尺寸最好是 16 的倍数因为经过四次下采样尺寸要能被 2 整除四次。如果输入是 572×572四次池化后是 35×35再上采样回来尺寸会对不上。所以实际使用中我通常把输入统一 resize 到 256×256 或 512×512省去很多麻烦。3.4 数据加载与预处理别让脏数据毁掉模型分割任务的数据组织和分类任务不一样。输入是一张图标签也是一张同尺寸的图每个像素的值代表类别。用 PyTorch 的 Dataset 和 DataLoader 来封装重点是保证图像和标签同步做变换。常见的预处理包括统一尺寸、归一化、随机翻转、随机旋转、颜色抖动。注意几何变换必须同时作用在图像和标签上否则标签就对不上了。归一化只对图像做标签保持整数类别。from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as T import torchvision.transforms.functional as F class SegDataset(Dataset): def __init__(self, img_paths, mask_paths, size256): self.img_paths img_paths self.mask_paths mask_paths self.size size def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(L) mask Image.open(self.mask_paths[idx]).convert(L) img F.resize(img, [self.size, self.size]) mask F.resize(mask, [self.size, self.size], interpolationT.InterpolationMode.NEAREST) img F.to_tensor(img) mask torch.from_numpy(np.array(mask)).long() return img, mask标签 resize 时一定要用最近邻插值否则会出现不存在的类别值。这个坑我踩过用双线性插值后标签里冒出 0.5 这种值训练直接报错。3.5 训练循环与损失函数落地训练循环的骨架和普通分类任务差不多前向传播、算损失、反向传播、更新参数。区别在于损失函数的选择和评估指标。我一般用 BCEWithLogitsLoss 加 Dice Loss 的组合。BCE 负责逐像素分类Dice 负责区域重叠。两者权重各 0.5 起步根据任务调整。优化器用 Adam学习率 1e-3 或 1e-4配合 ReduceLROnPlateau 在验证指标不提升时降学习率。def dice_loss(pred, target, eps1e-6): pred torch.sigmoid(pred) pred pred.view(-1) target target.view(-1).float() inter (pred * target).sum() return 1 - (2 * inter eps) / (pred.sum() target.sum() eps) criterion_bce nn.BCEWithLogitsLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, patience5)训练时每个 epoch 记录训练损失和验证 Dice保存验证指标最好的模型权重。不要只看训练损失分割任务很容易过拟合验证指标才是选模型的依据。4. 常见问题与排查技巧实录4.1 尺寸对不上、通道数报错怎么查这是新手最常遇到的问题。报错信息通常是 “size mismatch” 或 “channels mismatch”。排查思路很简单在 forward 里每一步打印张量形状看哪一步开始不对。尺寸对不上多半是输入不是 16 的倍数或者某次池化和上采样次数不匹配。通道数对不上多半是拼接时忘了算上采样后的通道数。比如 up4 输出 512 通道e4 也是 512 通道拼接后是 1024所以 dec4 的输入必须是 1024。这个数字关系要理清楚。提示写模型时先把通道数在纸上画一遍比在代码里反复试错快得多。4.2 训练损失不降、Dice 一直很低如果损失从一开始就不降先检查学习率是不是太大试试降到 1e-4。如果损失降但 Dice 不涨可能是类别极度不平衡Dice Loss 权重调高一些。如果训练集 Dice 高、验证集低那是过拟合加数据增强、加 Dropout、减小模型容量。还有一个隐蔽问题标签值域不对。二分类分割标签应该是 0 和 1如果标签是 0 和 255BCE 会算错。用torch.unique(mask)检查一下标签里到底有哪些值这个习惯能省很多时间。4.3 显存不够用的几种解法显存不够时按代价从低到高依次尝试减小 batch size、减小输入尺寸、把 DoubleConv 里的通道数整体减半、把转置卷积换成双线性插值、用混合精度训练。混合精度在 PyTorch 里用torch.cuda.amp就能开通常能省 30% 到 50% 显存速度也有提升。如果这些都不够那就只能上更小的模型或者多卡了。但大多数学习和中等规模项目前几招就够用了。4.4 边界分割毛糙、小目标漏检边界毛糙通常是跳跃连接没起到作用检查拼接是不是真的把浅层特征传过去了。小目标漏检则和损失函数、下采样倍数有关。如果目标本身只有几个像素四次下采样后可能就消失了。这时候可以减少下采样次数或者用空洞卷积替代部分池化保持分辨率的同时扩大感受野。我在一个细胞分割项目里遇到过类似问题最后把编码器改成三次下采样最底层用空洞卷积小目标召回率明显改善。所以 U-Net 不是固定不变的根据任务调整深度和感受野是常规操作。常见问题可能原因排查与解决尺寸 mismatch输入非 16 倍数resize 到 256 或 512通道 mismatch拼接通道算错画图核对每层通道数损失不降学习率过大降到 1e-4 重试Dice 低类别不平衡提高 Dice Loss 权重显存不足batch 或尺寸过大减 batch、开混合精度边界毛糙跳跃连接失效检查 concat 是否正确小目标漏检下采样过多减少池化、用空洞卷积4.5 几个让我少走弯路的实操习惯第一个习惯是先过拟合一个小数据集。拿 10 张图训练看能不能把训练 Dice 打到 0.99。如果能说明模型和数据管道没问题再上全量数据。如果不能问题一定在代码或数据里别急着调参。第二个习惯是可视化中间特征图。把编码器各层输出画出来看看浅层是不是保留了边界深层是不是有语义响应。这比盯着损失曲线有用得多。第三个习惯是固定随机种子。分割任务对初始化敏感固定种子后对比不同改动才有意义。torch.manual_seed(42)加上 numpy 和 random 的种子三行代码的事。第四个习惯是保存预测结果图。每个 epoch 存几张验证集的预测叠加图肉眼一看就知道模型在学什么。有时候指标涨了但图很难看说明指标和实际需求有偏差早点发现早点调整。5. 从能跑到好用几个进阶改造方向U-Net 跑通之后如果想进一步提升效果有几个方向值得试。一是把编码器换成预训练的主干网络比如 ResNet 或 EfficientNet利用大规模数据学到的特征小数据集上提升很明显。二是加入注意力机制在跳跃连接处做通道或空间注意力让网络自己决定哪些特征更重要。三是用深监督在解码器每一层都加辅助损失加速收敛。四是把损失换成 Tversky Loss 或 Focal Loss针对特定不平衡场景调优。这些改造不需要推翻重来都是在现有代码上加模块。我的建议是一次只改一个地方改完固定种子跑一遍对比验证指标确认有效再保留。分割任务的可复现性很重要东改西改最后不知道哪个起作用是很容易掉进去的坑。代码写到最后你会发现 U-Net 的魅力不在于它多复杂而在于它给了一个足够清晰、足够灵活的框架。理解了编码器、解码器、跳跃连接这三件事剩下的就是根据你的数据和任务去微调。我到现在做分割项目第一版基线还是 U-Net跑通了再想怎么改。这个习惯让我少了很多盲目尝试也更容易定位问题到底出在哪。 U-Net图像分割实战：从架构原理到PyTorch代码实现与调优
1. U-Net 为什么值得反复拆解第一次接触 U-Net 是在做医学影像分割的项目里当时手头只有几百张标注图像试过几个常规的编解码结构效果都不太理想。后来换成 U-Net在同样数据量下 Dice 系数直接上了一个台阶。从那时起我就意识到这个 2015 年提出的结构之所以到现在还被大量引用和改造不是因为它复杂恰恰是因为它把“少样本、高精度、强定位”这三件事用最朴素的方式捏在了一起。U-Net 本质上是一个编码器-解码器结构编码器负责逐层下采样、提取从纹理到语义的抽象特征解码器负责逐层上采样、把特征图恢复到原图分辨率最终输出逐像素的分类结果。它最标志性的设计是跳跃连接把编码器每一层的特征图直接拼接到解码器对应层让浅层的高分辨率细节不至于在下采样过程中被彻底丢掉。这个思路听起来简单但在分割任务里极其关键——没有它边界会糊成一团有了它即便训练样本很少网络也能靠这些“抄近路”的特征把轮廓抠得比较准。这篇文章面向的是想真正把 U-Net 跑起来的人不管你是刚学完卷积神经网络、想找一个完整项目练手的学生还是需要在医学影像、遥感、工业质检等场景里落地分割模型的工程师我都会从架构设计讲到 PyTorch 代码实现把每一步的“为什么”说清楚。我不会只贴一段能跑的代码就完事而是把参数选择、数据组织、损失函数、训练技巧、常见报错都摊开讲让你看完能自己改、自己调、自己排查。2. 网络架构与核心原理拆解2.1 整体结构一个对称的“U”形是怎么来的U-Net 的名字来自它的形状。把网络画出来左边一路下采样右边一路上采样中间用跳跃连接横向连起来整体像一个字母 U。这个形状不是审美选择而是功能决定的。左侧编码器通常由若干个“卷积块 下采样”组成。每个卷积块一般是两次 3×3 卷积每次后面接 ReLU再跟一个 2×2 最大池化把空间尺寸减半、通道数翻倍。以输入 572×572 的灰度图为例经过几次下采样后特征图会变成 28×28 甚至更小但通道数从 64 涨到 512语义信息越来越浓缩。右侧解码器则反过来先上采样把尺寸放大一倍然后和左侧对应层的特征图拼接再做两次 3×3 卷积。最后一层用 1×1 卷积把通道数压到类别数输出分割图。这里有个细节很多人第一次看会忽略原始 U-Net 论文里输入和输出尺寸并不相等因为当年用的是 valid 卷积每卷一次尺寸就缩一点。现在我们在 PyTorch 里实现时通常用 padding1 的 same 卷积让输入输出尺寸保持一致这样数据组织会简单很多。这是基于当前主流实践的合理调整不是对原论文的否定。2.2 跳跃连接U-Net 的灵魂所在如果只能保留 U-Net 的一个设计那一定是跳跃连接。它的作用可以从两个角度理解。从信息流角度看编码器在下采样时虽然感受野变大、语义变强但每个像素对应的空间位置信息被不断稀释。到了最底层一个特征点可能对应原图上很大一块区域边界早就模糊了。解码器如果只靠这些高层特征往上恢复分割结果就会像用低分辨率图片放大一样边缘全是锯齿和毛刺。跳跃连接把编码器浅层的高分辨率特征直接送过来相当于给解码器提供了一份“边界参考图”。从梯度角度看跳跃连接还是一条短路径。反向传播时梯度可以从解码器直接流回编码器浅层不需要穿过整个深层网络。这对训练稳定性有帮助尤其是在数据量不大、网络又比较深的时候能缓解梯度消失的问题。实际操作中拼接方式一般是通道维度上的 concatenation而不是相加。相加会丢失一部分信息拼接则让网络自己决定怎么融合。代价是通道数会翻倍显存占用增加所以如果你的显卡比较紧张可以在拼接后加一个 1×1 卷积把通道压回去这是常见的轻量化改法。2.3 上采样的几种做法与选择逻辑解码器里的上采样有好几种实现方式选哪种会直接影响效果和速度。最常用的是转置卷积也叫反卷积。它通过学习参数来放大特征图理论上更灵活但容易产生棋盘格伪影尤其是在步长和卷积核尺寸不匹配的时候。另一种是双线性插值 普通卷积先插值放大再卷积融合效果通常更平滑参数也更少。还有一种更简单的最近邻插值速度最快但比较粗糙一般只在轻量场景用。我在实际项目里的经验是如果追求精度且显存充足用转置卷积如果追求稳定和速度用双线性插值加卷积如果只是做原型验证最近邻也能凑合。没有绝对优劣关键看你的任务对边界精度的要求有多高。2.4 损失函数分割任务不能只看准确率分割任务有个特点背景像素往往远多于前景像素。比如一张医学影像里病灶可能只占几个百分点。如果只用普通的交叉熵损失网络会倾向于把所有像素都预测成背景准确率看起来很高但实际什么都没分出来。所以 U-Net 类任务里常用的损失是Dice Loss或者交叉熵 Dice 的组合。Dice 系数衡量的是预测区域和真实区域的重叠程度对类别不平衡不敏感。它的定义是 2×交集 / (预测面积 真实面积)取值在 0 到 1 之间越接近 1 越好。作为损失使用时通常写成 1 - Dice。另一个常见选择是BCEWithLogitsLoss它把 Sigmoid 和交叉熵合在一起数值稳定性更好。如果是多类别分割就用 CrossEntropyLoss。我的习惯是二分类分割用 BCE Dice 各占一半权重多分类用 CrossEntropy Dice这样既能保证像素级分类准确又能优化区域重叠。3. PyTorch 代码实现与关键细节3.1 环境准备把 PyTorch 装明白在写模型之前先把环境弄干净。我见过太多人卡在安装这一步后面代码没问题却跑不起来很打击积极性。如果你用 Anaconda推荐单独建一个环境不要和 base 混在一起。命令大概是先创建一个 Python 3.9 或 3.10 的环境然后根据你的显卡情况选择安装命令。有 NVIDIA 显卡且驱动正常的话去 PyTorch 官网查对应 CUDA 版本的安装命令用 conda 或 pip 都行。没有显卡就用 CPU 版本学习阶段完全够用只是训练慢一些。注意不要盲目复制网上的安装命令CUDA 版本、Python 版本、操作系统三者要匹配。装完后用torch.cuda.is_available()验证一下返回 True 才说明 GPU 可用。验证环境是否正常可以跑一段极简代码创建一个随机张量做一次卷积看输出尺寸对不对。这一步能排除大部分环境问题。3.2 卷积块与下采样模块的写法U-Net 的编码器由重复的卷积块组成所以先把它封装成一个类。每个卷积块包含两次 3×3 卷积每次后面接 BatchNorm 和 ReLU。BatchNorm 能加速收敛、稳定训练ReLU 提供非线性。padding 设为 1保证卷积后尺寸不变。下采样用 2×2 最大池化步长为 2这样尺寸正好减半。为什么用最大池化而不是平均池化因为最大池化保留的是最显著的特征响应在分割任务里更有利于突出边界和纹理。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.block(x)这段代码看起来简单但有几个点值得说。inplaceTrue能省一点显存但如果你后面要复用 ReLU 前的值就别用。BatchNorm 在 batch size 很小时表现不稳定如果显存只够跑 batch size 为 1 或 2可以考虑换成 GroupNorm 或 InstanceNorm这是医学影像分割里常见的调整。3.3 编码器、解码器与跳跃连接的组装编码器就是四个“卷积块 池化”的堆叠通道数依次是 64、128、256、512。每次池化后特征图尺寸减半。最底层再过一个卷积块通道数到 1024这是整个网络的“语义瓶颈”。解码器每一步先上采样然后把编码器对应层的特征拼过来再过一个卷积块。拼接时要注意通道数对齐上采样后的通道数加上编码器送来的通道数正好是卷积块的输入通道数。class UNet(nn.Module): def __init__(self, in_ch1, out_ch2): super().__init__() self.enc1 DoubleConv(in_ch, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.enc4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(512, 1024) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.dec4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)这段代码可以直接跑但要注意输入尺寸最好是 16 的倍数因为经过四次下采样尺寸要能被 2 整除四次。如果输入是 572×572四次池化后是 35×35再上采样回来尺寸会对不上。所以实际使用中我通常把输入统一 resize 到 256×256 或 512×512省去很多麻烦。3.4 数据加载与预处理别让脏数据毁掉模型分割任务的数据组织和分类任务不一样。输入是一张图标签也是一张同尺寸的图每个像素的值代表类别。用 PyTorch 的 Dataset 和 DataLoader 来封装重点是保证图像和标签同步做变换。常见的预处理包括统一尺寸、归一化、随机翻转、随机旋转、颜色抖动。注意几何变换必须同时作用在图像和标签上否则标签就对不上了。归一化只对图像做标签保持整数类别。from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as T import torchvision.transforms.functional as F class SegDataset(Dataset): def __init__(self, img_paths, mask_paths, size256): self.img_paths img_paths self.mask_paths mask_paths self.size size def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(L) mask Image.open(self.mask_paths[idx]).convert(L) img F.resize(img, [self.size, self.size]) mask F.resize(mask, [self.size, self.size], interpolationT.InterpolationMode.NEAREST) img F.to_tensor(img) mask torch.from_numpy(np.array(mask)).long() return img, mask标签 resize 时一定要用最近邻插值否则会出现不存在的类别值。这个坑我踩过用双线性插值后标签里冒出 0.5 这种值训练直接报错。3.5 训练循环与损失函数落地训练循环的骨架和普通分类任务差不多前向传播、算损失、反向传播、更新参数。区别在于损失函数的选择和评估指标。我一般用 BCEWithLogitsLoss 加 Dice Loss 的组合。BCE 负责逐像素分类Dice 负责区域重叠。两者权重各 0.5 起步根据任务调整。优化器用 Adam学习率 1e-3 或 1e-4配合 ReduceLROnPlateau 在验证指标不提升时降学习率。def dice_loss(pred, target, eps1e-6): pred torch.sigmoid(pred) pred pred.view(-1) target target.view(-1).float() inter (pred * target).sum() return 1 - (2 * inter eps) / (pred.sum() target.sum() eps) criterion_bce nn.BCEWithLogitsLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, patience5)训练时每个 epoch 记录训练损失和验证 Dice保存验证指标最好的模型权重。不要只看训练损失分割任务很容易过拟合验证指标才是选模型的依据。4. 常见问题与排查技巧实录4.1 尺寸对不上、通道数报错怎么查这是新手最常遇到的问题。报错信息通常是 “size mismatch” 或 “channels mismatch”。排查思路很简单在 forward 里每一步打印张量形状看哪一步开始不对。尺寸对不上多半是输入不是 16 的倍数或者某次池化和上采样次数不匹配。通道数对不上多半是拼接时忘了算上采样后的通道数。比如 up4 输出 512 通道e4 也是 512 通道拼接后是 1024所以 dec4 的输入必须是 1024。这个数字关系要理清楚。提示写模型时先把通道数在纸上画一遍比在代码里反复试错快得多。4.2 训练损失不降、Dice 一直很低如果损失从一开始就不降先检查学习率是不是太大试试降到 1e-4。如果损失降但 Dice 不涨可能是类别极度不平衡Dice Loss 权重调高一些。如果训练集 Dice 高、验证集低那是过拟合加数据增强、加 Dropout、减小模型容量。还有一个隐蔽问题标签值域不对。二分类分割标签应该是 0 和 1如果标签是 0 和 255BCE 会算错。用torch.unique(mask)检查一下标签里到底有哪些值这个习惯能省很多时间。4.3 显存不够用的几种解法显存不够时按代价从低到高依次尝试减小 batch size、减小输入尺寸、把 DoubleConv 里的通道数整体减半、把转置卷积换成双线性插值、用混合精度训练。混合精度在 PyTorch 里用torch.cuda.amp就能开通常能省 30% 到 50% 显存速度也有提升。如果这些都不够那就只能上更小的模型或者多卡了。但大多数学习和中等规模项目前几招就够用了。4.4 边界分割毛糙、小目标漏检边界毛糙通常是跳跃连接没起到作用检查拼接是不是真的把浅层特征传过去了。小目标漏检则和损失函数、下采样倍数有关。如果目标本身只有几个像素四次下采样后可能就消失了。这时候可以减少下采样次数或者用空洞卷积替代部分池化保持分辨率的同时扩大感受野。我在一个细胞分割项目里遇到过类似问题最后把编码器改成三次下采样最底层用空洞卷积小目标召回率明显改善。所以 U-Net 不是固定不变的根据任务调整深度和感受野是常规操作。常见问题可能原因排查与解决尺寸 mismatch输入非 16 倍数resize 到 256 或 512通道 mismatch拼接通道算错画图核对每层通道数损失不降学习率过大降到 1e-4 重试Dice 低类别不平衡提高 Dice Loss 权重显存不足batch 或尺寸过大减 batch、开混合精度边界毛糙跳跃连接失效检查 concat 是否正确小目标漏检下采样过多减少池化、用空洞卷积4.5 几个让我少走弯路的实操习惯第一个习惯是先过拟合一个小数据集。拿 10 张图训练看能不能把训练 Dice 打到 0.99。如果能说明模型和数据管道没问题再上全量数据。如果不能问题一定在代码或数据里别急着调参。第二个习惯是可视化中间特征图。把编码器各层输出画出来看看浅层是不是保留了边界深层是不是有语义响应。这比盯着损失曲线有用得多。第三个习惯是固定随机种子。分割任务对初始化敏感固定种子后对比不同改动才有意义。torch.manual_seed(42)加上 numpy 和 random 的种子三行代码的事。第四个习惯是保存预测结果图。每个 epoch 存几张验证集的预测叠加图肉眼一看就知道模型在学什么。有时候指标涨了但图很难看说明指标和实际需求有偏差早点发现早点调整。5. 从能跑到好用几个进阶改造方向U-Net 跑通之后如果想进一步提升效果有几个方向值得试。一是把编码器换成预训练的主干网络比如 ResNet 或 EfficientNet利用大规模数据学到的特征小数据集上提升很明显。二是加入注意力机制在跳跃连接处做通道或空间注意力让网络自己决定哪些特征更重要。三是用深监督在解码器每一层都加辅助损失加速收敛。四是把损失换成 Tversky Loss 或 Focal Loss针对特定不平衡场景调优。这些改造不需要推翻重来都是在现有代码上加模块。我的建议是一次只改一个地方改完固定种子跑一遍对比验证指标确认有效再保留。分割任务的可复现性很重要东改西改最后不知道哪个起作用是很容易掉进去的坑。代码写到最后你会发现 U-Net 的魅力不在于它多复杂而在于它给了一个足够清晰、足够灵活的框架。理解了编码器、解码器、跳跃连接这三件事剩下的就是根据你的数据和任务去微调。我到现在做分割项目第一版基线还是 U-Net跑通了再想怎么改。这个习惯让我少了很多盲目尝试也更容易定位问题到底出在哪。