深度学习里的分割模型一直是个绕不开的话题。尤其是U-net它从2015年提出到现在依然是医学影像、遥感图像、工业缺陷检测等领域里的常青树很多新模型做对比实验时第一个拿来比较的baseline就是它。这篇博文就和大家一起从零开始用PyTorch把U-net的网络结构搭出来并且讲清楚每一层设计背后的逻辑而不是简单地把代码贴出来就完事。文章的核心内容会覆盖U-net的编码器-解码器架构、跳跃连接的关键思路、数据增强的必要性以及一个能直接跑起来的训练与验证流程。我会尽量把“为什么要这样设计”也讲透比如为什么用跳跃连接而不是简单的残差相加为什么数据增强在分割任务里特别重要。适合刚开始做图像分割、或者已经跑通分类任务想转向像素级任务的读者参考对PyTorch有一定基础但还没系统搭过分割网络的开发者也会很有帮助。1. U-net网络的核心设计思路1.1 U形结构的本质多尺度特征融合U-net的名字来源于它独特的U形对称结构。左边是一路编码器逐层下采样分辨率越来越小通道数越来越多右边是解码器逐层上采样分辨率逐步恢复到原始输入尺寸。这种结构本质上是在做一件很关键的事情多尺度特征提取。编码器部分其实就是卷积神经网络的堆叠每一层通过池化或步长为2的卷积把特征图缩小一半同时增加通道数。这个过程有点像我们在看一幅大图时先看整体轮廓再看局部细节逐级细化。解码器则相反通过上采样把低分辨率的特征图慢慢放大回原始尺寸。这里有个非常重要的设计点U-net在每个尺度的解码器部分会把对应编码器同一层的特征图拼接concat过来这就是所谓的跳跃连接。我经常把跳跃连接比作是“带着轮廓画细描边”。编码器前面的层保留了很多边缘、纹理之类的底层细节信息这些信息在后期的深层网络里早就被抽象化了。通过跳跃连接这些高频细节直接补充到解码过程里帮助网络恢复目标的精确边界。和ResNet里那种把特征图相加的残差连接相比U-net选择在通道维度上做拼接通道数翻倍信息量更大。这是两类连接的本质区别我觉得这是理解U-net首先要站在心里的一个点。1.2 医疗影像场景下的先天优势U-net最早是为医学细胞图像分割设计的它最明显的优点是即使在训练样本很少的情况下也能获得不错的分割效果。这一点对医疗影像非常重要因为标注一张高质量的病变区域图需要专业医生花费大量时间普通的公开数据集动辄几百例都已经算多的了。为什么U-net在少样本场景下还能站得住脚关键在于编码器-解码器加跳跃连接的结构本身对参数的利用效率比单纯的分类网络更高。U-net在每次下采样之前都连续做两次3x3卷积这使得网络在每一层都能提取到充分的局部特征而不是像一些深层分类网络那样靠极深的堆叠换取性能。再加上数据增强的辅助少样本下的鲁棒性会进一步提升。除了医学影像U-net在遥感图像中的道路提取、建筑物识别工业领域的表面缺陷分割自动驾驶场景里的车道线与人车目标分割里也被大量使用。可以说U-net是很多分割任务的第一步把这套结构吃透再看DeepLab、SegNet或者Transformer类的分割模型时会容易得多。1.3 为什么选择PyTorch复现PyTorch近年几乎是学术研究和工程落地里最主流的深度学习框架之一。它的动态计算图让网络结构搭建和调试显得非常透明尤其是当我们需要在forward函数里做跳跃连接这种跨层拼接操作时直接写一行代码就能完成不用像静态图那样需要考虑图结构的编译细节。另外PyTorch的生态也越来越成熟预训练模型库、各类数据Loader、分布式训练、ONNX导出都支持得不错。对初学者来说它有一点像在写普通的Python代码而不是在使用另一个平台。所以我要用PyTorch来搭建U-net这个过程本身就是学习框架特性的好途径。环境准备方面Anaconda PyTorch CUDA的组合是最常用的后文会有清晰的环境搭建指引。2. PyTorch环境搭建与项目结构准备2.1 基于Anaconda的GPU环境配置搭建环境是很多初学者最容易卡住的第一步。建议先安装Anaconda它会帮我们管理好Python和常用的科学计算包。然后创建一个独立的conda环境避免不同项目之间的依赖冲突。我们用命令行来操作具体步骤是conda create -n unet python3.8 conda activate unet进入环境后安装PyTorch。这里要注意选择对应CUDA版本的安装命令可以去PyTorch官网获取最新的命令。一般优先选择GPU版本因为训练分割网络时即便输入图像不大计算量也非常可观。示例如下pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118安装完成后可以用一段小代码验证GPU是否可用import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU detected)如果输出True并且能看到显卡型号说明环境就绪。2.2 项目目录与依赖组织一个良好的项目结构能让后续开发和调试省下很多时间。我的做法是将数据、模型、训练脚本、工具函数分开放置。比如unet-project/ ├── data/ │ ├── images/ │ └── masks/ ├── src/ │ ├── model.py │ ├── dataset.py │ ├── train.py │ └── utils.py └── checkpoints/模型定义放在model.py里数据集读取放在dataset.py里训练逻辑在train.py中一些可视化与指标计算的辅助函数放进utils.py。这样分离后每次想改动某个环节只需要进入对应的文件而不必在一个大脚本里翻来翻去。对于训练分割网络还需要额外安装一些科学计算与图像库例如OpenCV、NumPy、tensorboard以及用于图像增强的albumentations。这些都直接用pip安装即可。数据增强库我后面会重点讲它对分割任务非常重要。2.3 数据准备的基本要求做分割任务数据和标注是核心。每张输入图像需要对应一张和它大小一致的标注图标注图的每个像素代表一个类别。例如在医学图像中通常背景为0目标区域为1。如果有多类则为0, 1, 2等整数编码。自己整理数据时要注意图像与标注图的对齐关系。有时候标注图是PNG格式的索引图长宽必须和原图完全一致。可以在数据读取阶段加一个断言assert来检查尺寸如果不一致就打印错误信息避免训练时出现张量形状不匹配的尴尬问题。另外建议将数据封装成PyTorch的Dataset和DataLoader这样在训练时可以方便地做shuffle、多进程读取和batch组装。3. 完整U-net结构搭建与PyTorch代码实现3.1 编码器模块的实现细节U-net的编码器由多个block组成。每个block包含两次3x3卷积每次卷积后接ReLU激活函数。在block的最后进行一次步长为2的最大池化或者步长为2的卷积来完成下采样。经典论文中使用的是最大池化而现代实现里也有人用步长为2的卷积替代两者的作用都是让特征图尺寸减半。我在复现时用的是两个卷积加一个最大池化的组合。每次卷积操作会将输入图像的通道数翻倍第一个block比较特殊直接从输入通道数变到第一个特征通道数。下面是编码器模块的代码示例import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, kernel_size3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)这里加了BatchNorm这是我做的一点改动。原论文年代比较早没有加BN但现在大家普遍发现加BN之后训练更稳定收敛也更快。注意卷积的padding设置为1这样卷积操作后特征图尺寸不改变尺寸只靠池化或上采样来控制。3.2 通过跳跃连接完整的U-net主体构建完整的U-net结构包括四个下采样阶段的编码器、一个中间的瓶颈层以及四个上采样阶段的解码器。在每个解码器block中首先对上采样的特征图做转置卷积或者上采样后接卷积然后和对应编码器block的输出特征图在通道维度上拼接。拼接的代码非常直观decoder_input torch.cat([encoder_output, decoder_output], dim1)dim1表示在通道维度上拼接。拼接之后接着再通过一个DoubleConv来融合两类特征。每个解码器层输出的特征图通道数会减半直到最后一层输出一个通道数等于类别数的特征图。完整网络的一个可能实现如下我用一个列表演示编码器与解码器的嵌套关系并保留跳跃连接的中间结果class UNet(nn.Module): def __init__(self, in_channels3, num_classes1, features[64, 128, 256, 512]): super(UNet, self).__init__() self.downs nn.ModuleList() self.ups nn.ModuleList() self.pool nn.MaxPool2d(kernel_size2, stride2) # Encoder for feature in features: self.downs.append(DoubleConv(in_channels, feature)) in_channels feature # Bottleneck self.bottleneck DoubleConv(features[-1], features[-1]*2) # Decoder for feature in reversed(features): self.ups.append( nn.ConvTranspose2d(feature*2, feature, kernel_size2, stride2) ) self.ups.append(DoubleConv(feature*2, feature)) self.final_conv nn.Conv2d(features[0], num_classes, kernel_size1) def forward(self, x): skip_connections [] for down in self.downs: x down(x) skip_connections.append(x) x self.pool(x) x self.bottleneck(x) skip_connections skip_connections[::-1] for idx in range(0, len(self.ups), 2): x self.ups[idx](x) skip skip_connections[idx//2] if x.shape ! skip.shape: x nn.functional.interpolate(x, sizeskip.shape[2:]) x torch.cat((skip, x), dim1) x self.ups[idx1](x) return self.final_conv(x)这段代码里的核心点在于把编码器每个block输出的特征图存进skip_connections列表然后反转顺序让最浅层的特征图在最后一步上采样时被拼接。如果你真的理解了这一行“反转”你就掌握了U-net结构的一半精髓。3.3 输入输出尺寸计算与多类别适配在医学分割中一张常见的输入图可能是512x512或256x256。由于每次池化会让尺寸减半经过四次下采样后特征图变成输入尺寸的1/16。比如输入是512池化后依次为256、128、64、32经过瓶颈层后还是32。解码阶段逐步上采样回64、128、256、512因此输出图尺寸和输入图完全一致。这种“尺寸不变”的设计对工程实现很友好因为不需要为不同的输入尺寸设计复杂的调整逻辑。唯一需要注意的是如果输入图像的宽高不是16的倍数会在跳跃连接拼接时出现尺寸不一致。我在代码中加了一个nn.functional.interpolate将上采样后的特征图插值到和跳跃连接相同的尺寸算是一个防御性的处理。如果分割任务有多个类别比如背景、目标A、目标B三类只需要把final_conv的输出通道数设为3即可。在训练损失上使用CrossEntropyLoss。如果是二分类分割输出通道设为1配合Sigmoid和BCEWithLogitsLoss使用。4. 训练循环与损失函数、指标分析4.1 分割任务里损失函数怎么选U-net的训练离不开合适的损失函数。最常见的是二分类场景下的BCEWithLogitsLoss它对Logits直接计算二值交叉熵数值上更稳定。多分类场景用CrossEntropyLoss。但分割任务有一个特殊的痛点前景区域在整幅图中占比往往很小普通交叉熵会倾向于把像素全预测为背景导致指标虚高但实际效果很差。所以实际项目里我一般会使用组合损失。Dice Loss是一种基于区域重叠度的损失函数它直接优化预测区域和真实区域的Jaccard相似程度对前景背景不敏感。也常常把Dice Loss和交叉熵相加共同作为监督信号。数学上Dice系数的计算公式是[ Dice \frac{2 \times |X \cap Y| smooth}{|X| |Y| smooth} ]这里smooth是一个防止分母为零的平滑项一般取1或1e-5。PyTorch的代码实现也不复杂只需要把预测和标签都转成浮点数逐元素相乘后求和即可。Dice Loss就是1减去Dice系数。把两者按一定权重加起来训练效果往往出人意料地好。4.2 优化器、学习率和训练步长设置我在训练U-net时最常用的优化器是Adam初始学习率设置在1e-3到3e-4之间。相比SGDAdam更容易上手对学习率的敏感度也相对低一些。如果追求更极致的性能也可以换成SGD配合动量但需要更多调参时间。有一个很实用的小技巧是使用学习率衰减策略。比如每训练N个epoch后学习率乘上0.9或0.5或者使用PyTorch里现成的ReduceLROnPlateau在指标停滞时自动降低学习率。分割网络通常训练轮数较多配合衰减策略往往能在最后几轮进一步提升Dice分数。训练步长方面batch size的选择主要看显卡显存。一个包含64个通道、尺寸为512x512的特征图就会占掉很多显存所以batch size设为2或4都很常见。如果显存不够常见的方法是减小batch size、缩小输入尺寸或者开启梯度累积来模拟更大的batch size。4.3 训练循环与模型保存的最佳实践训练循环的结构对熟悉PyTorch的人来说并不陌生准备数据、前向传播、计算损失、反向传播、优化器更新、验证、保存模型。但在分割任务里验证环节不能只看loss还需要额外计算Dice系数或IoU。下面是一个精简的训练循环骨架for epoch in range(epochs): model.train() train_loss 0.0 for images, masks in train_loader: images, masks images.to(device), masks.to(device) outputs model(images) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() * images.size(0) model.eval() valid_dice 0.0 with torch.no_grad(): for images, masks in valid_loader: images, masks images.to(device), masks.to(device) outputs model(images) valid_dice compute_dice(outputs, masks).item() * images.size(0)保存模型时我推荐把epoch、模型权重、优化器状态、最佳指标等都打包存入一个字典使用torch.save保存为.pth文件。这样后期恢复训练或做模型推理时可以一次性加载所有状态不用手动去拼凑。torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_dice: best_dice, }, checkpoints/unet_epoch{}.pth.format(epoch))4.4 使用TensorBoard可视化训练过程训练分割网络光看终端里的数字是不够的。我习惯用TensorBoard记录loss曲线和Dice曲线在每个epoch结束后记录数值还可以把一些输入图像和预测结果直接画到TensorBoard里直观查看每一轮的变化。PyTorch的torch.utils.tensorboard提供了简明的接口。在训练循环里每隔一定步数或每个epoch把损失和指标写入SummaryWriter。展示预测结果时注意把模型输出的概率图通过Sigmoid转成0到1的范围再转成绘图能接受的NumPy数组效果如下图示意此处无图仅做方法演示。可视化步骤里最容易踩的坑是图像尺寸或数据类型的转换。TensorBoard的add_image接口要求传入的tensor是HxWxC或者NxCxHxW格式且数值范围是0到255或0到1。我通常将预测结果转为uint8的三通道图再来写TensorBoard。这样一张图里就能同时看到原图、真实标注和预测结果调试效率高很多。5. 数据增强与过拟合应对策略5.1 分割任务的数据增强核心逻辑分割任务经常面临训练数据量不足的问题数据增强就是最直接的应对方法。和普通分类任务不同分割任务在做几何增强时除了要变换输入图像还必须要同步变换对应的标签图否则标签会和图像内容错位训练出来必出问题。常用的增强方式包括水平翻转、垂直翻转、随机旋转、随机缩放、随机裁剪、颜色抖动等。如果使用albumentations库事情会非常简单因为它提供了同时处理image和mask的接口。例如import albumentations as A from albumentations.pytorch import ToTensorV2 transform A.Compose([ A.RandomRotate90(p0.5), A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.2), A.Resize(512, 512), ToTensorV2(), ]) augmented transform(imageimage, maskmask) aug_image augmented[image] aug_mask augmented[mask]使用albumentations库时最好一次性把图像和mask都传给transform它可以自动保持空间变换的一致性。这一点非常关键自己手工写随机翻转时如果忘记同步处理mask模型训练指标会非常奇怪。5.2 处理类别不平衡问题图像分割里前景区域过小、背景区域巨大是常态。假设一张512x512的图像中病变区域只占几万像素那网络很容易学成“全预测为背景”因为这样的准确率可能也有99%以上。如果只看准确率你会觉得模型已经非常优秀但画出预测图来会发现一片漆黑。面对这类问题Dice Loss本身就比纯交叉熵更稳健。为了进一步缓解不平衡还可以使用带权重的交叉熵给前景像素更高的权重。另外对损失函数做组合也是一种思路比如在Dice Loss为主的基础上加一点交叉熵辅助稳定梯度。在训练时多观察一下验证集上的Dice与IoU不要过于依赖loss数值或准确率。5.3 防止过拟合的常见手段分割网络参数规模并不小如果数据量只有几十张很容易在训练集上快速过拟合。除了数据增强我还会用在验证集指标连续多个epoch不再提升时降低学习率并保存当前最好的模型权重这个习惯帮我避免了很多“眼看要起飞结果飞歪了”的情况。另一个经验是不要一开始就用全尺寸原图。如果显卡不够或者数据量少先缩放到256x256或320x320训练跑通流程后再往上加分辨率。这样可以快速验证数据管道有没有问题避免浪费大量时间在超高分辨率上调试。等模型结构没大问题后再切换到更高的分辨率做最终训练。6. 常见报错与调试排查实录6.1 尺寸不匹配与Skip Connection报错这是复现U-net时最容易遇到的报错。有一种典型情况是在跳跃连接环节编码器输出的特征图尺寸和解码器上采样后的尺寸差了一个像素。原因是输入图像的宽高不是2的整数次幂经过两次池化后多出或少了半个像素。如果像我前面那样在代码里加入interpolate做尺寸对齐这个问题就能迎刃而解。假如你的报错信息类似“size mismatch for skip_connection”第一反应应该是去检查各层特征图尺寸。最简单的排查方法是在forward函数里加一个print(tensor.shape)打印每个关键tensor的尺寸然后和理论推算对比。不要盲目改卷积参数先确认尺寸变化是否符合预期。6.2 显存溢出CUDA Out of Memory训练分割网络时显存溢出非常常见尤其是使用高分辨率输入时。一个非常高效的办法是在训练脚本里设置torch.cuda.empty_cache()在每轮验证前清理缓存。如果还是溢出优先把batch size降到1或者把训练图像尺寸从512降到384或者256。我试过很多次在U-net里面把512的图降到256显存占用可以减少接近一半。另一个值得留意的点是PyTorch在反向传播时会保存中间激活值用于计算梯度这是显存消耗的大头。如果显存捉襟见肘可以使用混合精度训练AMP也就是用torch.cuda.amp模块把前向传播中一部分计算改为半精度。这项技术在V100、A100或RTX 30系以上显卡上支持得很好开启后显存占用和训练速度都会有明显改善。6.3 收敛不稳定与Loss消失如果训练过程中loss出现NaN或者一开始就停在某个值附近不动最常见的原因有两个一是学习率过大二是标签值错位。我遇到过一次很隐蔽的问题数据集里的mask是RGB格式的彩色图而不是索引图结果在计算损失时被当成类别ID来用loss直接爆炸。后来我统一把mask读取为单通道灰度图只取其中某个颜色通道作为标签值问题才解决。学习率方面如果loss在初始阶段就产生剧烈震荡通常就是学习率太高了。可以把学习率降到3e-5再试。如果loss下降非常缓慢则可能学习率太低。另外使用BCEWithLogitsLoss时模型的输出不需要先经过Sigmoid因为该函数内部已经组合了Sigmoid和交叉熵计算这一点尤其容易和手动加Sigmoid后接BCELoss混淆。6.4 常见问题速查表下面这个表格是我在实际操作中积累的快速排查笔记遇到问题时可以对照检查现象可能原因解决建议Loss为NaN学习率过大或标签值异常降低学习率检查mask是否为0到类别数减1的整数训练精度高但预测效果差类别极度不平衡改用Dice Loss或为交叉熵增加权重GPU显存不足输入分辨率过大或batch过大减小输入尺寸降低batch size开启混合精度训练验证集Dice不上升数据增强过强或学习率不匹配调低增强概率调整初始学习率增加衰减策略跳跃连接尺寸对不上输入尺寸不是16倍数使用interpolate统一尺寸或预处理时resize到合适大小7. 训练结果评估与可视化输出7.1 用Dice和IoU客观评估模型语义分割的常用评价指标是Dice系数和IoUIntersection over Union。IoU的计算方式是两个集合的交集面积除以并集面积反映了预测区域和真实区域的重叠程度。Dice和IoU之间存在数学关系两者趋势通常一致但Dice对小型区域的波动更敏感所以很多医疗影像任务习惯报Dice。评估时有一点要特别注意必须在验证集或测试集上计算而不是训练集。训练集上的指标会因为模型看到了数据而偏高不能真实反映泛化能力。另外如果任务是多类别分割最好分别计算每个类别的Dice和IoU然后取平均而不是把所有类别混合在一起算一个总指标。这样才能发现某个类别是不是完全没被模型识别出来。7.2 预测结果可视化与保存模型训练完成后效果到底行不行可视化预测结果是最直观的。可以把测试集里几张图像的预测结果保存为原图、真实mask、预测mask的三联图。这样贴到报告或者论文里时审稿人一眼就能看到分割质量。保存预测mask时要注意PyTorch模型的输出如果是多通道特征图需要先做Softmax或Sigmoid然后取每个像素概率最大的类别作为预测类别。然后我们一般会把预测结果乘上一个最大像素值255再保存为图像文件。如果只用plt.imshow直接显示浮点矩阵颜色会显得很奇怪因为Matplotlib默认会把浮点数归一化到0到1之间。我自己就踩过这个坑输出图看起来全是黑的后来才发现只是显示范围的问题。8. 对U-net结构的深度思考与改进心得8.1 U-net的局限性和常见变体虽然U-net效果很扎实但它也有明显的局限性。最直接的短板是整个网络属于全卷积结构没有全局建模能力。卷积核的感受野毕竟是有限的对于非常依赖上下文信息的长距离分割任务比如大范围病灶区域的分割U-net可能在边界连续性上表现一般。近几年出现了一批U-net变体试图解决这个问题。比如UNet引入了密集跳跃连接重构了解码器接收多级编码器特征的路径Attention U-net在跳跃连接前加入注意力门控模块让模型自动关注重要的区域再比如TransUNet把Transformer放在U-net的编码器末端借助自注意力机制增强全局特征建模。这些变体在公开数据集上的表现通常有一定提升但结构复杂度也上升了。我在实际项目中用Attention U-net比较多注意力模块的参数量不算大但确实能改善一些目标边缘模糊的情况。对于绝大多数任务原生U-net仍然是一个非常可靠的起点。先跑通原版再按需引入变体这才是稳妥的路线。8.2 替换卷积模块的实践心得如果想让U-net在特定场景下更强有一个低成本的方式把核心的DoubleConv模块替换为ResBlock或DenseBlock。我试过把普通卷积换成残差卷积后深层网络在接近20层以上的训练稳定性更好梯度回传也更顺畅。这个改动在代码上只是再写一个新的卷积模块在原来的DoubleConv里加一个捷径分支然后保证输入输出通道数一致就可以。换用空洞卷积是另一种思路。在不降低特征图分辨率的前提下增大感受野是处理一些细小或长条形目标的好办法。比如在瓶颈层使用空洞率分别为2、4、8的空洞卷积可以在不增加参数量的情况下看到更大的上下文。针对裂缝检测这类任务这个改动往往能带来肉眼可见的边界准确率提升。8.3 从网络结构到系统工程U-net的搭建只是分割项目的第一步。如果要把它落地到实际生产环境还需要考虑推理速度、模型压缩、部署推理框架等多个环节。比如用TensorRT对ONNX模型做加速或者在边缘设备上使用量化后的模型推理速度能提升数倍。很多同学做项目时只关注网络结构本身但工程里往往推理性能和分割精度同样重要。我的建议是在做U-net复现的阶段先不要引入太花哨的结构把基本流程理清楚数据准备、模型搭建、训练、验证、可视化、保存导出。等这条流水线顺畅之后再去优化网络的某个局部或者做工程加速。这样后续做任何改进时都能有一个稳定可靠的baseline作为对照不会被各种变量干扰。9. 完整实操步骤回顾与小技巧分享9.1 从零到训练成功的一份简明流程清单我整理了一份适合第一次跑通U-net的流程清单照着走一般不会有太大问题安装Anaconda和PyTorch GPU环境验证CUDA可用。准备图像和mask数据保持文件名一一对应尺寸统一。实现Dataset类在__getitem__中同步读取图像与mask。定义U-net模型先小尺寸快速验证forward输出尺寸正确。选择合适的损失函数建议BCEWithLogitsLoss或Dice Loss组合。初始化Adam优化器和学习率衰减策略。运行一个最简单的epoch观察显存占用和loss是否正常。多epoch训练每次验证后保存最佳权重。加载最佳权重可视化预测结果并与真实mask对比。计算验证集Dice和IoU评估模型稳定性。9.2 训练时几个容易忽略的细节细节一别忘了调用model.train()和model.eval()。如果忘记切换到eval模式BatchNorm层在验证时仍然会使用当前batch的统计量导致验证结果不稳定。这个坑相当隐蔽因为loss和指标看起来都能正常计算但数值可能异常偏高或者忽高忽低。细节二在验证阶段要使用torch.no_grad()包裹前向传播否则会构建计算图浪费显存和算力。这个问题在batch比较小、网络比较深的时候会放大有时候验证做一半就内存溢出了。细节三DataLoader的num_workers参数不要设置得太大Windows系统下如果设为0通常能避开很多奇怪的并行bug。Linux系统下可以设置为4或8来加速数据读取。如果数据加载变成瓶颈还有一种办法是用albumentations做在线增强时适当把图像尺寸调小换取更高的吞吐量。9.3 效果不佳时优先调整什么很多初学者拿到一个不理想的结果第一反应就是去改网络结构或者换更大的模型。但根据我的经验90%的U-net效果不佳问题并不在结构而在数据或者训练配置。优先检查顺序应该依次是是否有标签错位、是否数据增强过度、是否类别不均衡、是否学习率不合适最后再考虑模型本身。如果训练集和验证集上的loss都居高不下先怀疑数据通道顺序有问题或者mask没有正确转换到tensor。如果训练集loss很低但验证集指标差这是典型的过拟合需要增加数据增强强度或加入正则化手段。如果训练和验证的loss都在缓慢下降只是绝对值高多半是任务本身的像素级回归有难度可以试着换个损失函数或引入更合理的边界权值。10. 写在最后的实验心得把U-net用PyTorch完整搭一遍对我来说不只是复习了一次经典网络结构更像是一次对分割任务全流程的重新梳理。从数据怎么读到训练如何稳再到模型怎么保存和评价每一个环节都有很多容易踩的坑。我印象最深的是第一次跑分割任务时因为忘记切换model.eval()验证Dice一直抖动不停排查了很久才发现是这么基础的原因。从那以后我每次训练循环都会把train/eval切换和no_grad包裹放在显眼的位置反复检查。如果你正在复现U-net我建议在动手写代码前先手动画一遍整个网络的数据流图标好每层输出张量的尺寸和通道数。这个看起来笨拙的练习能帮你省下大量debug时间。等你亲手把U-net跑起来看到预测mask和真实mask逐渐重合的那一刻就会明白这套结构的“优雅”到底在哪里。后面再去学Transformer分割模型或者注意力机制也会有一个非常踏实的落脚点。