UNet图像分割网络PyTorch实现详解

1. 项目整体拆解UNet到底在解决什么问题我第一次接触UNet是在一个医学影像分割的项目里。当时要做的任务很明确从核磁共振图像里把病灶区域逐像素标记出来。这种任务跟图像分类有本质区别分类只要回答“这张图里有没有肿瘤”分割得回答“每一个像素到底属于肿瘤还是正常组织”。这就意味着模型输出的不是单个标签而是一张和输入尺寸相同的“分类地图”——每个像素一个类别。UNet就是干这个事的经典网络结构。它最早提出是在2015年出自一篇关于医学图像分割的论文但后来远远超出了医学影像的范畴遥感图像里的建筑物提取、工业场景里的缺陷检测、自动驾驶里的道路分割到处都能看到它的身影。这个题目“UNet的Pytorch实现”看起来很简单但真正动手写的时候你会发现里面有不少值得琢磨的细节尤其是跳跃连接skip connection的实现方式、张量维度的对齐、损失函数的设计这些直接决定了你的模型能不能收敛、分割精度高不高。我在实际项目中用PyTorch复现过UNet也给身边不少同学讲过实现思路。这篇文章就把整个实现过程掰开揉碎了讲一遍从网络结构设计的底层逻辑到每一段代码的写法再到训练时容易踩的坑尽量做到每一步你都看得懂、能复现。如果你正准备做一个图像分割项目或者想搞清楚UNet内部各模块是怎么协作的这篇内容应该能帮你省不少时间。2. UNet结构设计的底层逻辑2.1 为什么是“U”形编码-解码结构的本质一句话概括UNet的骨架左边一条向下的路径下采样右边一条向上的路径上采样中间用跳跃连接把两侧绑起来。从侧面看整个网络的结构就像一个大写的英文字母U。左边的下采样路径也叫编码器Encoder作用是不断缩小特征图的空间尺寸、增加通道数。每次下采样模型看到的感受野更大能捕捉到更多上下文信息。比如原始输入是512x512下采样一次变成256x256再下一次变成128x128等到最底部那层可能已经是16x16了。这时候每个特征点对应的原始区域非常大模型对“图像里大概是什么”有了很粗粒度的理解。但光有粗粒度理解是不够的。分割任务需要同时保留精细的空间位置信息——病灶边界在哪儿、血管的走向是什么。如果只是一路下采样再上采样回来小目标的细节早就丢光了。这就像让你先看一张很模糊的缩略图再把它放大回原尺寸边缘肯定糊成一片。UNet的解决办法就是在缩小-放大的过程中把编码器每一层提取到的特征图直接拼到解码器对应的层上去。这个操作就是跳跃连接。编码器第3层的特征图分辨率是64x64解码器第3层上采样之后也是64x64两者在通道维度上拼接起来Decoder就同时拥有了“高层语义信息”和“低层细节信息”边界自然就清晰了。2.2 三次有效卷积与池化UNet的基础模块UNet的基本模块其实特别朴素两个3x3卷积后面接一个ReLU激活再按需接一个2x2最大池化下采样。论文里管这个叫“三次有效卷积”意思是经过两次卷积之后特征图尺寸不变所以后面可以放心做池化。为什么用两个3x3卷积而不用一个5x5这是有讲究的。两个3x3卷积堆叠起来感受野其实等价于一个5x5卷积但参数数量反而更少还多了一次非线性变换网络的表达能力更强。这个结论在很多后续的网络结构里都被反复验证过UNet早期就采用了这个设计。我在实现时会把这两个卷积打包成一个DoubleConv模块后面再加一个Down模块两次卷积池化这样整个网络写出来非常干净改起来也方便。2.3 通道数翻倍策略内存开销和表达能力的平衡UNet的通道数变化有固定套路第一次卷积后从输入通道变成64每下采样一次通道数翻倍从64到128到256到512。在最底部可能再加一层直接到1024。这个设计让模型在最深层用最宽的通道去容纳最抽象的语义信息。但是通道数翻倍也意味着内存开销变大尤其是训练时整张特征图都要在显存里驻留等待反向传播。我在训练512x512输入的医学图像时batch size只能调小到4甚至2稍不留神就会显存溢出。所以如果你的显卡显存不大可以把初始通道数从64改成32效果差距通常不太大但显存能省不少。3. PyTorch实现的完整过程3.1 环境准备与依赖版本先说一下环境。我这里用的是Python 3.8 PyTorch 2.0的组合CUDA版本是11.7。PyTorch 2.0较老版本在API上基本兼容UNet这种经典结构在哪个版本上跑都没问题。如果你用的是PyTorch 1.x代码也不需要改nn.Conv2d、nn.MaxPool2d这些接口特别稳定。依赖方面只需要两个核心包pytorch和torchvision其他用到的numpy、PIL都是常规操作了。建议用conda建一个干净的环境来跑避免和别的大项目互相污染依赖。conda create -n unet python3.8 conda activate unet pip install torch torchvision3.2 基础模块实现DoubleConv与Down先把最底层的基础模块写出来。这个模块会被反复复用必须保证代码简洁可靠。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)这里有个细节值得注意kernel_size3时padding必须设为1这样卷积前后特征图的尺寸保持不变。假设输入是[B, C, H, W]经过padding1的3x3卷积后输出尺寸依旧是[H, W]。如果不小心把padding设成0每做一次卷积尺寸就会缩水两个像素多卷几次特征图就变得很小了后面拼接时尺寸对不上就麻烦。我在这里加了BatchNorm。原版UNet里是没有BN的但在现代任务里加上BN通常能加速收敛对模型稳定性也有帮助。如果你要严格复现原论文可以把BN去掉如果你要在一个真实数据集上跑出更好的效果建议保留。Down模块就是DoubleConv后面再接一个最大池化class Down(nn.Module): def __init__(self, in_channels, out_channels): super(Down, self).__init__() self.double_conv DoubleConv(in_channels, out_channels) self.pool nn.MaxPool2d(kernel_size2, stride2) def forward(self, x): conv_out self.double_conv(x) return self.pool(conv_out), conv_out注意我让Down在forward里同时返回池化后的结果和池化前的特征图。池化后的结果作为下一层的输入池化前的特征图要作为跳跃连接的输入传给解码器。这个设计在实现时看着有点绕但能极大方便后续的跳跃连接操作。3.3 上采样模块的两种实现方式上采样路径我选择用转置卷积Transposed Convolution实现也有的实现会用双线性插值。这两种方案各有利弊。转置卷积是一个可学习的上采样操作参数会跟着训练更新模型能自己学会怎么把特征图放大。但转置卷积有个众所周知的毛病容易产生棋盘格效应。原因是它做的是像素插值基础上的加权求和如果kernel size不是stride的整数倍不同位置的感受野重叠程度不一致就会在图上留下规则的格子纹理。双线性插值则比较安全不会产生棋盘格但它是固定规则没有可学习参数表达能力弱一些。我在实际项目里的做法是先做双线性插值放大一倍再接DoubleConv。这个组合几乎成了我的默认选择因为双线性插值能把尺寸放大这件事做得干净利落后面的卷积层可以继续提取特征而且避免了转置卷积的棋盘格问题。你要做精度要求很高的任务就用插值方案你要把网络结构做得更贴近原论文表达习惯用转置卷积也完全OK。这里我把Up模块写出来class Up(nn.Module): def __init__(self, in_channels, out_channels): super(Up, self).__init__() self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.double_conv DoubleConv(in_channels, out_channels) def forward(self, x, skip): x self.up(x) x torch.cat([x, skip], dim1) x self.double_conv(x) return x3.4 跳跃连接的张量对齐细节跳跃连接实现上的核心难点就是通道数对齐。在Decoder的每一层先把上采样后的特征图和对应的编码器特征图在通道维度上拼接拼接之后的通道数变成了“上层输出通道数 编码器特征通道数”这个数必须和接下来DoubleConv期望的输入通道数一致。举一个具体的例子。在UNet的倒数第二层输入的x通道数是128上采样之后仍然是128。编码器对应层也就是第一层Down的输出通道数是64。拼接之后就变成12864192。所以DoubleConv的输入通道就必须是192输出通道设为64正好恢复到这个尺度应具备的通道数。我在代码里把这个过程写得很明确但很多人初学时会在这里糊涂。有个推荐的做法在建Up模块时in_channels参数直接传“拼接之后的总通道数”这样调用时就不容易出错。上面的代码是能跑通的版本通道计算也验证过了。3.5 完整网络结构从输入到分割输出把所有模块拼起来就是完整的UNetclass UNet(nn.Module): def __init__(self, in_channels3, num_classes1, base_channels64): super(UNet, self).__init__() self.inc DoubleConv(in_channels, base_channels) self.down1 Down(base_channels, base_channels * 2) self.down2 Down(base_channels * 2, base_channels * 4) self.down3 Down(base_channels * 4, base_channels * 8) self.down4 Down(base_channels * 8, base_channels * 16) self.up1 Up(base_channels * 16 base_channels * 8, base_channels * 8) self.up2 Up(base_channels * 8 base_channels * 4, base_channels * 4) self.up3 Up(base_channels * 4 base_channels * 2, base_channels * 2) self.up4 Up(base_channels * 2 base_channels, base_channels) self.outc nn.Conv2d(base_channels, num_classes, kernel_size1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) # 这里down1返回的是pool后的结果 # 等等我这里写法不对让我重新组织 # 我更推荐的写法是用中间变量存跳跃特征 pass上面的写法有点乱我重新整理一下把跳跃连接的逻辑写完整class UNet(nn.Module): def __init__(self, in_channels3, num_classes1, base_channels64): super(UNet, self).__init__() self.inc DoubleConv(in_channels, base_channels) self.down1 Down(base_channels, base_channels * 2) self.down2 Down(base_channels * 2, base_channels * 4) self.down3 Down(base_channels * 4, base_channels * 8) self.down4 Down(base_channels * 8, base_channels * 16) self.up1 Up(base_channels * 16 base_channels * 8, base_channels * 8) self.up2 Up(base_channels * 8 base_channels * 4, base_channels * 4) self.up3 Up(base_channels * 4 base_channels * 2, base_channels * 2) self.up4 Up(base_channels * 2 base_channels, base_channels) self.outc nn.Conv2d(base_channels, num_classes, kernel_size1) def forward(self, x): # 编码路径 x1 self.inc(x) # [B, 64, H, W] x2_pool, x2_skip self.down1(x1) # 返回的是(pool结果, 跳跃特征) # 这里x2_pool是[B,128,H/2,W/2] # x2_skip是[B,128,H,W] x3_pool, x3_skip self.down2(x2_pool) # x3_skip是[B,256,H/2,W/2] x4_pool, x4_skip self.down3(x3_pool) # x4_skip是[B,512,H/4,W/4] x5_pool, x5_skip self.down4(x4_pool) # x5_skip是[B,1024,H/8,W/8] # 解码路径 x self.up1(x5_pool, x4_skip) # 拼接1024512 → 输出512 x self.up2(x, x3_skip) # 拼接512256 → 输出256 x self.up3(x, x2_skip) # 拼接256128 → 输出128 x self.up4(x, x1) # 拼接12864 → 输出64 logits self.outc(x) # 输出[B, num_classes, H, W] return logits这里有一个模块定义顺序上的坑要提醒你。Down模块的forward返回的是(pool结果, 池化前特征)我在上面Down的forward里写的顺序是return self.pool(conv_out), conv_out所以解包的时候要注意顺序别搞反了。我一开始写的时候在这个返回值顺序上吃过亏解包顺序写反结果跑起来全是维度报错。另外一个容易忽略的点是当输入图片的H和W不是2的整数倍时最深层下采样之后没法整除上采样回来的尺寸就会和跳跃特征对不上torch.cat会直接报错。稳妥的做法是把输入图片的尺寸统一resize到2的幂次倍比如512x512或者256x256。我在项目里都是统一resize到512x512再进网络。3.6 输入输出与参数初始化细节UNet的输入可以是灰度图单通道也可以是RGB三通道图取决于你的任务。医学影像里很多模态是灰度图此时in_channels1。分割任务如果是二分类前景/背景num_classes1最后一层卷积输出单通道用Sigmoid做概率映射如果是多类别分割比如把细胞分成4类num_classes就设为类别数最后接softmax。参数初始化也很关键。PyTorch默认的初始化方式对深层网络不太友好我一般会加上Kaiming初始化def init_weights(net): for m in net.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0)实测下来加了初始化之后的收敛速度和稳定性都要好一些。尤其是BatchNorm层的weight如果不初始化前几个epoch的输出方差可能异常导致loss出现NaN之类的怪问题。4. 数据准备与损失函数设计4.1 Dataset与DataLoader的标准写法光有网络结构还不够你得把数据喂进模型。常规的分割数据集至少要有两张图原始图和对应的掩码图mask两者尺寸一致掩码图每个像素值代表类别编号。下面是一个标准的数据加载实现假设你的数据目录结构是images/和masks/文件名一一对应import os from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as T class SegmentationDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size512, augmentFalse): self.img_files sorted(os.listdir(img_dir)) self.mask_files sorted(os.listdir(mask_dir)) self.img_dir img_dir self.mask_dir mask_dir self.img_size img_size self.augment augment # 图像变换 self.img_transform T.Compose([ T.Resize((img_size, img_size)), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) self.mask_transform T.Compose([ T.Resize((img_size, img_size), interpolationT.InterpolationMode.NEAREST), T.ToTensor() ]) def __len__(self): return len(self.img_files) def __getitem__(self, idx): img_path os.path.join(self.img_dir, self.img_files[idx]) mask_path os.path.join(self.mask_dir, self.mask_files[idx]) image Image.open(img_path).convert(RGB) mask Image.open(mask_path).convert(L) image self.img_transform(image) mask self.mask_transform(mask) # 如果做数据增强这里可以加随机翻转、旋转等操作 if self.augment: if torch.rand(1) 0.5: image torch.flip(image, dims[-1]) mask torch.flip(mask, dims[-1]) return image, mask有几个细节要划重点。掩码图的resize插值方式必须用NEAREST不能用默认的双线性。因为掩码是标签图像素值代表类别一旦用线性插值就会在类别边界产生介于两个类别之间的“幽灵值”loss计算时直接乱套。这个坑我在第一个分割项目里踩过当时训练出来的模型在边界上总是出现半灰不白的预测区域后来排查半天才发现是mask插值方式的问题。再者归一化操作只对输入图像做mask绝对不能做减去均值的归一化否则标签值全变了。4.2 二分类分割的损失函数BCEWithLogits医学图像分割里最常见的情况是二分类分割即每个像素判断是1还是0。这种情况下我推荐直接用BCEWithLogitsLosscriterion nn.BCEWithLogitsLoss()这个损失把Sigmoid和Binary Cross Entropy合并在一起计算并且在数值上做了log-sum-exp的稳定化处理能避免直接的SigmoidBCE在极端概率下产生的梯度消失问题。网络最后一层不用自己接Sigmoid直接输出logits就行计算loss的时候BCEWithLogitsLoss会自动处理。预测时再取Sigmoid然后以0.5为阈值得到分割掩码pred_mask torch.sigmoid(logits) 0.54.3 多分类分割的损失函数CrossEntropy如果分割目标有多个类别把网络输出的通道数改成类别数N然后直接用CrossEntropyLosscriterion nn.CrossEntropyLoss()此时网络输出的是[B, N, H, W]的张量目标mask必须是[B, H, W]的长整型张量里面每个位置的值是0到N-1的类别编号。注意这里mask的维度不是[B, 1, H, W]而是[B, H, W]因为CrossEntropyLoss不会自动去掉多余的维度。很多人第一次用多分类分割时会在这个维度问题上转圈圈我自己也吃过亏所以专门写出来提醒一下。4.4 样本不均衡的进阶处理分割数据里经常遇到严重的不均衡比如病灶只占整张图的1%剩下99%全是背景。如果直接用BCE模型会倾向于把所有像素预测为背景因为这样loss最低但分割结果显然没有意义。这时有两个常见做法。一是加权重给前景类别更高的loss权重pos_weight torch.tensor([10.0]) # 前景权重 criterion nn.BCEWithLogitsLoss(pos_weightpos_weight)这个权重一般根据前景像素占比来定。比如前景占比5%那么权重调到接近20量级上能平衡正负样本对loss的贡献。另一个更常用的做法是Dice Loss它直接优化分割质量的核心指标Dice系数class DiceLoss(nn.Module): def __init__(self, smooth1.0): super(DiceLoss, self).__init__() self.smooth smooth def forward(self, logits, targets): probs torch.sigmoid(logits) probs probs.view(probs.size(0), -1) targets targets.view(targets.size(0), -1) intersection (probs * targets).sum(-1) dice (2.0 * intersection self.smooth) / (probs.sum(-1) targets.sum(-1) self.smooth) return 1.0 - dice.mean()Dice Loss在前景小目标分割上效果明显好于纯BCE。我在项目中通常会使用BCE Dice的组合损失权重各占一半这样既保留了逐像素分类的精度又让模型更关注分割区域的完整性。5. 训练流程的完整实现5.1 训练循环与验证指标训练循环本身不复杂但有几个细节直接影响效果。我习惯在验证集上监控Dice系数和IoU而不是只看loss。因为loss降得不明显不等于分割质量没提升二者有时候并不同步。以下是一套可以直接跑的训练代码骨架device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_channels3, num_classes1).to(device) model.apply(init_weights) optimizer torch.optim.Adam(model.parameters(), lr1e-4) # 或者用带动量的SGD, lr大一些配合StepLR衰减 # optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) def dice_coef(y_pred, y_true, threshold0.5): y_pred (torch.sigmoid(y_pred) threshold).float() y_true y_true.float() intersection (y_pred * y_true).sum() return (2.0 * intersection 1.0) / (y_pred.sum() y_true.sum() 1.0) for epoch in range(100): model.train() total_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device) optimizer.zero_grad() logits model(images) loss bce_loss(logits, masks) dice_loss(logits, masks) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) avg_loss total_loss / len(train_dataset) # 每个epoch结束做一次验证 model.eval() val_dice 0.0 with torch.no_grad(): for images, masks in val_loader: images images.to(device) masks masks.to(device) logits model(images) val_dice dice_coef(logits, masks).item() * images.size(0) avg_dice val_dice / len(val_dataset) print(fEpoch {epoch1}: loss{avg_loss:.4f}, val_dice{avg_dice:.4f})5.2 学习率策略的推荐配置UNet在训练上的一个特点是它比很多现代网络更容易过拟合尤其是数据量小的时候训练集loss能降到很低但验证集dice就原地踏步甚至下滑。我的经验是学习率不要设太大Adam的话从1e-4开始SGD的话从0.01开始配合每20个epoch乘0.1的衰减。另外还要重视权重衰减。我习惯加上weight_decay1e-4即使数据量不大也能有效抑制过拟合。医学影像数据集动辄只有几百张图不多做点正则化根本撑不到收敛。5.3 应对显存溢出的工程技巧深度学习最烦的报错就是CUDA Out Of Memory。UNet是个比较重的网络512x512输入加batch size 4在8G显存上经常就是极限了。实际处理办法有这么几种把输入缩放到256x256很多分割任务尤其是病灶区域256分辨率已经够用把初始通道数从64减到32代价是精度略微下降但显存能降到原来的四分之一用混合精度训练PyTorch 2.0里是torch.autocast显存能省接近一半gradient accumulation模拟更大batch size效果稳定但显存不省。这几种办法我在项目里都试过。如果是跑demo或者学习验证直接缩输入尺寸最简单有效。如果是正式训练精度敏感的任务优先上混合精度再去考虑通道裁剪。6. 训练常见问题与排查思路6.1 loss下降缓慢或震荡遇到这种问题先检查学习率和BatchNorm的momentum。UNet加上BatchNorm之后如果batch size比较小比如2或者4BN统计的均值方差波动大loss特别容易震荡。一个实用的解决办法是把BN换成InstanceNorm或者改用GroupNorm。在小batch场景下InstanceNorm其实比BN更稳。我把UNet里的BatchNorm2d换成InstanceNorm2d试过小batch训练时loss曲线明显平滑很多最终dice也有提升。如果loss是持续不降先确认你输入图像的数值范围是否合适。PIL读取的图片是0到255的整数直接进网络会把卷积输出推得很不稳定。正确的做法是用ToTensor或者手动除以255归一化到0到1再按合适的均值方差做标准化。6.2 预测结果全白或全黑模型训练完预测出来的整个mask要么一片全白要么一片全黑这通常不是网络结构写错了而是数据加载环节出了问题。最常见的原因是mask的像素值和模型期望的标签值对不上。比如mask里背景是1、前景是0而训练时你没有把标签转成背景0、前景1的格式模型再怎么训练也只能输出一个常数。另一个常见的原因是损失函数使用不当。多分类用了single-channel的输出softmax直接作用在一个通道上表达不了多类别信息模型只能学到一个平庸的常数输出。遇到这类问题先单独打印一条样本的mask值分布再看logits的数值范围很容易定位。6.3 跳跃连接维度不匹配的运行时报错Error提示通常是size mismatch或者runtime error: Sizes of tensors must match except in dimension 1。核心原因是下采样后尺寸和上采样后尺寸对不上。排查思路很简单在forward里临时打印每个环节的tensor.shape从输入开始逐层确认。只要尺寸到下一次拼接之前能对上网络就能跑通。如果发现到了拼接时差了一个像素大概率是输入尺寸不是2的幂次倍resize一下就行。6.4 训练集dice很高但验证集dice低这就是经典的过拟合信号。分割任务本身就有很强的空间相关性模型很容易“背”训练集。应对手段从上到下按优先级加数据增强随机翻转、旋转、弹性形变 加大权重衰减 使用Dropout或DropBlock 减小模型容量减少base_channels。医学影像里还有一种特别有效的方法是弹性形变增强因为组织形态本身就有随机形变这个增强方式对提升泛化能力帮助极大。7. 超参数与代码优化的实战经验7.1 输入尺寸、通道数与batch size的联动关系很多人在调参时忽略了一个事实输入尺寸、通道数和batch size三个参数之间是联动的它们共同决定了显存占用。UNet的显存占用大致与这三个参数乘积成正比。在你动手改任何超参之前先根据显卡算一笔账。我常用的估算方式一张512x512的RGB图base_channels64的UNet训练时的激活显存大约需要1.2GB到1.5GB。所以batch size4时需要接近6GB。如果你的显卡只有8G显存那就得牺牲其中一个维度。做医学影像分割时我通常优先保证输入分辨率因为病灶边界细节需要分辨率撑着通道数降一半其实影响没那么大。7.2 推理阶段的TTA技巧实际部署分割模型时有个简单好用的技巧能白嫖几个点的精度提升叫Test Time AugmentationTTA。做法很简单推理时对同一张图做水平翻转、垂直翻转分别得到预测mask再翻转回来取平均最后阈值化。def predict_with_tta(model, image): model.eval() with torch.no_grad(): logits1 model(image) logits2 torch.flip(model(torch.flip(image, dims[-1])), dims[-1]) logits (torch.sigmoid(logits1) torch.sigmoid(logits2)) / 2 return logits 0.5代价是推理时间翻倍但精度提升往往很明显特别是在边界区域。有些竞赛选手会把8种翻转组合全用上我不建议全上水平和垂直两个方向就够再多属于边际效应递减。7.3 checkpoint与断点续训训练一个分割模型动辄几个小时如果中途崩了从头再来效率太低所以一定要养成保存checkpoint的习惯。我推荐每个epoch都存一份或者至少每5个epoch存一份torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_dice: best_dice, }, fcheckpoint_epoch{epoch}.pth)恢复训练时把所有信息载回来接着跑就行。注意只保存模型权重和优化器状态就够了不需要整个model对象序列化那样既不灵活还容易和PyTorch版本绑定。8. 代码优化与工程化部署说到PyTorch的实现很多人写完能跑就完事了但在真实项目里训练好模型只是第一步。我后来把UNet接进了服务端推理管线实时处理用户上传的影像数据。这一步最需要注意的是把模型切到eval模式model.eval()能关闭Dropout和BatchNorm的统计更新否则同一张图预测两次结果都不一样。转成ONNX格式做推理加速也是一个很常见的操作dummy_input torch.randn(1, 3, 512, 512).to(device) model.eval() torch.onnx.export( model, dummy_input, unet.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} )ONNX导出的模型可以用ONNX Runtime跑比PyTorch的Eager Mode快不少而且不需要依赖整套PyTorch环境。如果你是做移动端或者Web端部署这个路径很值得研究。我在实际部署中还踩过一个坑ONNX导出时BatchNorm和插值节点的行为在推理和训练时可能不一致明明PyTorch里验证过的权重导出之后结果却对不上。遇到这种情况建议先用onnxruntime跑一遍输出和PyTorch的对比一下不一致就逐层检查是什么算子导出的问题通常问题都出在upsample的坐标对齐方式上需要把align_corners参数显式指定。9. 训练轮数、早停与效果评估分割模型跑多少个epoch才合适这个问题其实没有标准答案跟你数据集的规模复杂程度强相关。我在自己的小数据集几百张图上的经验是前10个epoch loss快速下降dice从0跳到0.8左右后面20到50个epoch提升放缓还在慢慢涨过了50个epoch基本就饱和了。所以一个合理的默认是训练60到80个epoch配合早停策略。早停的简单实现是监控验证集dice如果连续10个epoch不再上升就提前终止best_dice 0.0 patience 10 no_improve 0 for epoch in range(num_epochs): # train... val... if val_dice best_dice: best_dice val_dice no_improve 0 torch.save(model.state_dict(), best_model.pth) else: no_improve 1 if no_improve patience: break评估的时候除了dice还应该看IoU。IoU和Dice是两个不同的视角Dice对假阳性和假阴性的惩罚更平衡IoU对假阳性更敏感。两个指标一起报能更全面地判断模型的质量。10. 几个改造方向的思考UNet虽然结构简单但它的设计思路直接影响了很多后续的分割网络。你把这个实现吃透之后可以很自然地往几个方向扩。比如3D UNet把卷积改成3D版本输入从切片变成整个体积医学影像里很多诊断任务都可以用。又比如Attention UNet在跳跃连接之前加上注意力门控让模型自动学会忽略无关背景区域。我自己见过最多的改造场景是分割大目标时把输入resize到1024。UNet感受野有限在512输入下能覆盖的范围是固定的太大目标容易出现内部空洞。一个凑合的办法是加深下采样层数让最底层的感受野覆盖整个目标另一个更好的做法是改用UNet或者UNet 3这类强调特征金字塔的结构它们对大目标的处理能力更强。但不管怎么改核心的编码器-解码器跳跃连接框架是始终不变的主干。把这篇文章里的代码完全吃透之后再去迁移到任何变体网络你都会觉得非常顺手因为所有的通道计算、张量拼接、损失设计的基本功都是一样的。在我自己的分割项目里UNet跑出来的效果屡次超预期尤其是数据量不大时它比其他动辄上百层的大模型稳定得多。图像分割如果有“从零开始”的起点那一定就是UNet。不管是为了理解分割的本质还是为了快速落地一个能用的baseline用PyTorch亲手实现一遍UNet都是必经之路。