UNet的Pytorch实现——别被这个标题吓到这其实是我最近手把手复现UNet时踩坑踩出来的完整记录。UNet在医学影像分割、卫星遥感图像处理、自动驾驶语义分割这些场景里几乎是入门必练的骨架网络Pytorch又是现在做深度学习实验最顺手的框架。这篇文章我会从网络结构的核心设计讲起重点拆解编码器、解码器、跳跃连接这些关键模块为什么这么搭再给出可以直接跑的完整代码包括数据流是怎么在这一堆卷积和拼接之间穿行的。无论是刚啃完一堆深度学习理论想动手练练的新手还是项目里需要快速嵌入一个分割模型的老手这篇都能省下你不少自己翻源码、查报错的时间。我复现这个项目的动机很实在网上关于UNet的代码一抓一大把但很多要么是古董写法要么只贴了核心类不贴训练流程对新手极不友好。真正要把UNet跑起来从模型定义、数据加载到损失函数、训练循环每一步都有细节坑稍不注意Loss就不收敛或者图像尺寸对不上直接爆维度错误。所以我决定自己完整地走一遍把能用、能改、能跑的方案记录下来。1. 为什么UNet是分割任务里的常青树1.1 UNet这个名字的来历UNet最早是2015年提出来的医学图像分割网络论文名字叫U-Net: Convolutional Networks for Biomedical Image Segmentation。它结构画出来是个大写字母U的形状左边是收缩路径右边是扩张路径中间底部是瓶颈层。这个设计在当时并不是什么惊天动地的创新卷积、池化、上采样、跳跃连接这些组件全都是现成的但作者把它们组合出一种非常高效的模式以至于到今天UNet及其变体依然是许多分割任务的首选baseline。它解决的核心问题很朴素图像分割需要同时知道这个东西在哪和这个东西是什么。分类网络比如VGG、ResNet特征是越往后越抽象空间位置信息被层层池化稀释最后的feature map往往只有原图的1/16甚至1/32拿去做分割就是知道是猫但不知道猫在画面的左上角还是右下角。UNet的思路是解码过程中把编码器每一层的高分辨率特征直接拼过来用跳跃连接把丢掉的空间细节补回来这样既能享受深层特征的语义能力又保住浅层特征的边缘细节。这里有个关键点值得展开跳跃连接不是简单的相加而是通道维度上的拼接。也就是说假设编码器某一层输出是256通道解码器对应层上采样后是128通道拼接后变成384通道再经过卷积压缩到128通道。这个先拼后压的操作让网络有机会自己学习怎么融合语义特征和空间特征比直接相加灵活得多。这也是UNet后来这么多变体里最常见的保留设计。1.2 为什么选择Pytorch而不是其他框架我在这个项目里选Pytorch有几个很实际的原因。首先Pytorch的调试体验确实好动态图机制让我可以随时print中间张量的形状出问题当场就能定位不需要像静态图那样先编译再执行。其次Pytorch社区在医学影像、视觉分割方向积累了大量现成代码和预训练权重复现论文、找参考都方便。再一个Pytorch对自定义数据集和自定义损失函数非常友好只需要继承Dataset类和nn.Module类就行整个流程心智负担小。有人可能会问TensorFlow的Keras接口不也很快吗如果只是调几行现成的模型Keras或许更快但一旦要做深度定制比如改个损失函数、加个自定义的注意力模块Keras的封装反而成了阻碍。Pytorch是一切皆张量、一切皆Module没有额外的抽象层理解起来更直接。对于学习UNet这种结构稍微复杂的网络用Pytorch动手撸一遍要比在黑盒里调参学到的多得多。2. 环境准备与数据预处理2.1 环境依赖与硬件建议动手之前先把环境搭明白。我这个项目用的是Pytorch 2.x版本Python 3.10显卡是某款中端卡显存8GB。如果你的机器没有GPU纯CPU也能跑只是训练速度慢很多建议把图像尺寸和batch size调小一点不然一个epoch能等半小时。以下是requirements.txt的关键内容torch2.0 torchvision0.15 numpy1.24 opencv-python4.8 Pillow10.0 tqdm4.65 matplotlib3.7这里有个特别需要提醒的坑Pytorch版本差异会导致某些API行为不同比如torch.nn.functional.interpolate的align_corners参数在不同版本里默认值不一样。如果你跟着老教程写代码上采样尺寸差一个像素是家常便饭。我这个项目的代码是在Pytorch 2.0以上版本测试通过的如果你用的版本比较老建议先升级。2.2 数据集准备一张灰度图对应一张标签UNet的标准数据组织方式是每张原始图像对应一张同样尺寸的标签掩码图标签图上每个像素的值代表类别编号。以最常见的二分类分割为例前景是1背景是0标签图就是一个二值图。如果是多类别分割标签图里就是0、1、2、3这样的整数每个数字代表一个类别。数据量方面很多公开分割数据集动辄几千上万张图但UNet的一个优势恰恰是小数据集也能训练得不错。如果自己标注几十张到一两百张图也够出一个能看的效果关键是做好数据增强。数据增强这块我一开始没太注意后面发现这其实是分割任务里最能提升泛化能力的一环。常用的增强手段包括随机水平翻转、随机垂直翻转、随机旋转、随机裁剪、亮度对比度扰动。跟分类任务不同分割任务的增强必须是图像标签同步变换否则标签和图像对不上训练等于白做。我自己实现了一个简单的同步增强函数import random import numpy as np import torch from torchvision import transforms from PIL import Image class SegmentationDataset(torch.utils.data.Dataset): def __init__(self, image_paths, mask_paths, image_size(256, 256)): self.image_paths image_paths self.mask_paths mask_paths self.image_size image_size def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image Image.open(self.image_paths[idx]).convert(RGB) mask Image.open(self.mask_paths[idx]).convert(L) # 缩放到统一尺寸 image image.resize(self.image_size, Image.BILINEAR) mask mask.resize(self.image_size, Image.NEAREST) image np.array(image, dtypenp.float32) / 255.0 mask np.array(mask, dtypenp.int64) # 随机水平翻转 if random.random() 0.5: image image[:, ::-1, :] mask mask[:, ::-1] # 随机垂直翻转 if random.random() 0.5: image image[::-1, :, :] mask mask[::-1, :] # HWC - CHW image torch.from_numpy(image.transpose(2, 0, 1)) mask torch.from_numpy(mask) return image, mask这段代码里有几个细节得说明。mask用NEAREST插值缩放是有讲究的因为mask是类别标签用双线性插值会把0和1混合出小数导致标签不再是整数训练直接报错。图像用BILINEAR则是为了保持图像平滑。如果你用OpenCV做增强cv2.flip有个参数需要注意水平翻转传1垂直翻转传0别弄反了。2.3 验证集划分与可视化检查数据准备好之后别急着开训。先画几张图看看图像和标签是否对齐这是我最强调的步骤因为标注工具导出的名字对不上、图片通道顺序不对、mask被不小心压缩成jpg导致出现伪边缘这些问题在训练前用肉眼看图全都能暴露到了训练里就变成loss下降极慢或者评测指标一直上不去的玄学。我的划分方式比较简单粗暴90%训练10%验证随机打乱。很多分割任务不需要严格的K折交叉验证因为这个任务本来就容易过拟合留一个验证集看曲线就够了。另外如果训练集和验证集来自不同批次的标注可能会有分布差异最好先打乱再划分避免按目录顺序导致验证集全是某个特定类型的样本。3. UNet核心代码实现与模块拆解3.1 编码器从图像到语义的逐层压缩UNet的编码器部分就是一个标准的卷积特征提取栈。我的实现里每一层用两次3x3卷积加ReLU激活然后接一个2x2最大池化做下采样。为什么用两次卷积而不是一次这是UNet的一个经典细节两次卷积让每个层级的特征提取更充分非线性表达能力更强。3x3卷积加padding1能保持特征图尺寸不变这样编码器的每个层级输出尺寸只由池化决定结构清晰稳定。以下是编码器模块的代码import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super(DoubleConv, self).__init__() self.double_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.double_conv(x)我把BatchNorm加进去了这是对原版UNet的一个常见改进。原版UNet没有BN层在那个年代不觉得有太大问题但实际操作中BN能显著加速收敛还能对初始化不那么敏感。有一点需要注意如果你的图像尺寸非常小比如32x32BN在小batch上可能不稳定可以换成InstanceNorm。我这里用的batch size是16图像256x256BN效果挺好。通道数的设计遵循典型UNet模式第一层64然后128、256、512瓶颈层1024。这个倍率选择跟显存直接相关如果显存吃紧可以把初始通道改成32或者16感受野和整体结构不受影响只是网络容量变小。3.2 解码器与跳跃连接把细节拼回来解码器部分要做的两件事上采样恢复分辨率拼接编码器的特征。我使用的是转置卷积上采样2x2转置卷积步长为2把特征图尺寸翻倍。也有人用插值上采样加1x1卷积的方式效果差不太多但转置卷积的参数是可学习的理论上能更灵活。实际测试中插值上采样的泛化更稳定而转置卷积如果学习率没调好容易在特征图出现棋盘格伪影。我这里保留转置卷积方案但在训练策略里加了学习率衰减来压制这个问题。class UpBlock(nn.Module): def __init__(self, in_channels, skip_channels, out_channels): super(UpBlock, self).__init__() self.up nn.ConvTranspose2d(in_channels, out_channels, kernel_size2, stride2) self.conv DoubleConv(out_channels skip_channels, out_channels) def forward(self, x, skip): x self.up(x) # 处理边缘尺寸不一致 if x.size(2) ! skip.size(2) or x.size(3) ! skip.size(3): diff_h skip.size(2) - x.size(2) diff_w skip.size(3) - x.size(3) x torch.nn.functional.pad(x, [diff_w // 2, diff_w - diff_w // 2, diff_h // 2, diff_h - diff_h // 2]) x torch.cat([x, skip], dim1) return self.conv(x)跳跃连接的核心就在这里——torch.cat把上采样后的深层特征和编码器的浅层特征在通道维度上拼起来。为什么这么简单的一个操作效果如此惊人我的理解是浅层特征保存着大量细节纹理和边缘信息深层特征包含着语义类别信息两者拼接后卷积层就有机会同时参考这两类信息去做决策。如果只用深层特征分割结果的边缘一定会很粗糙只用浅层特征模型就不理解自己分割的是什么。3.3 完整模型组装与输入输出尺寸验证把编码器和解码器拼起来再加上最后的1x1卷积分类层class UNet(nn.Module): def __init__(self, in_channels3, num_classes1): super(UNet, self).__init__() # 编码器 self.enc1 DoubleConv(in_channels, 64) self.pool1 nn.MaxPool2d(2) self.enc2 DoubleConv(64, 128) self.pool2 nn.MaxPool2d(2) self.enc3 DoubleConv(128, 256) self.pool3 nn.MaxPool2d(2) self.enc4 DoubleConv(256, 512) self.pool4 nn.MaxPool2d(2) # 瓶颈层 self.bottleneck DoubleConv(512, 1024) # 解码器 self.up1 UpBlock(1024, 512, 512) self.up2 UpBlock(512, 256, 256) self.up3 UpBlock(256, 128, 128) self.up4 UpBlock(128, 64, 64) # 输出层 self.output nn.Conv2d(64, num_classes, kernel_size1) def forward(self, x): # 编码器路径保存每层输出用于跳跃连接 skip1 self.enc1(x) x self.pool1(skip1) skip2 self.enc2(x) x self.pool2(skip2) skip3 self.enc3(x) x self.pool3(skip3) skip4 self.enc4(x) x self.pool4(skip4) # 瓶颈 x self.bottleneck(x) # 解码器路径逐层与跳跃连接拼接 x self.up1(x, skip4) x self.up2(x, skip3) x self.up3(x, skip2) x self.up4(x, skip1) return self.output(x)写完后第一件事绝对是拿一个假张量验证尺寸model UNet(in_channels3, num_classes1) x torch.randn(1, 3, 256, 256) y model(x) print(y.shape) # 期望输出 torch.Size([1, 1, 256, 256])注意原版UNet要求输入尺寸能被2整除四次也就是至少是16的倍数否则编码器的池化会导致尺寸向下取整最后解码器拼不上。如果输入图像是512x512或者1024x1024完全没问题。如果非要用不规则尺寸我在UpBlock里加了pad兜底但最稳妥的做法还是统一resize到256x256或512x512。3.4 损失函数与评估指标分割任务二分类场景最常用的损失就是BCEWithLogitsLoss把模型最后的输出直接喂进去内部先算sigmoid再算交叉熵。使用BCEWithLogitsLoss的好处是数值稳定性更好它通过log-sum-exp的技巧避免了sigmoid后梯度过小的问题。多分类场景则用CrossEntropyLoss它内部会自动做softmax。我实际用的损失是BCEWithLogitsLoss和Dice损失的组合。Dice损失直接优化分割结果与标注的像素级重叠率对前景背景比例极度不平衡的情况非常有效。医学图像里经常出现前景只占整张图不到5%的情况纯用BCE会导致模型倾向于全预测背景因为那样损失也不大。Dice损失的存在让模型不得不认真找出那个小小的前景。def dice_loss(pred, target, smooth1.0): pred torch.sigmoid(pred) pred pred.view(pred.size(0), -1) target target.view(target.size(0), -1) intersection (pred * target).sum(dim1) union pred.sum(dim1) target.sum(dim1) dice (2.0 * intersection smooth) / (union smooth) return 1.0 - dice.mean() # 训练时组合使用 bce nn.BCEWithLogitsLoss() loss bce(pred, target) dice_loss(pred, target)评估指标方面最直观的是IoU和Dice系数。训练过程中我每隔几个epoch就在验证集上算一次IoU把数字打出来看趋势。IoU的计算需要把网络输出通过sigmoid拿到概率值再按0.5阈值转成二值图跟标签对比计算交集除以并集。4. 训练循环与推理全流程4.1 训练循环监控Loss和验证指标训练循环本身不复杂但有几个地方如果写得不严谨会直接影响模型效果。我给出一个带验证与检查点保存的完整训练脚本骨架import torch import torch.optim as optim from tqdm import tqdm def train_one_epoch(model, dataloader, optimizer, criterion, device): model.train() total_loss 0.0 for images, masks in tqdm(dataloader): images images.to(device) masks masks.to(device).float().unsqueeze(1) preds model(images) loss criterion(preds, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader) def validate(model, dataloader, device): model.eval() dice_total 0.0 with torch.no_grad(): for images, masks in dataloader: images images.to(device) masks masks.to(device).float().unsqueeze(1) preds model(images) dice_total (1 - dice_loss(preds, masks)).item() return dice_total / len(dataloader) # 主训练配置 device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_channels3, num_classes1).to(device) optimizer optim.Adam(model.parameters(), lr1e-4) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemax, factor0.5, patience5) criterion torch.nn.BCEWithLogitsLoss() num_epochs 100 best_dice 0.0 for epoch in range(num_epochs): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) val_dice validate(model, val_loader, device) scheduler.step(val_dice) print(fEpoch {epoch1}/{num_epochs} | Train Loss: {train_loss:.4f} | Val Dice: {val_dice:.4f}) if val_dice best_dice: best_dice val_dice torch.save(model.state_dict(), best_unet.pth)这里有个经验点想分享optimizer我用Adamlr设1e-4起步然后配合ReduceLROnPlateau动态调整。如果训练初期loss完全不动先把lr调到1e-3试如果还是不动检查数据归一化和mask的dtype。另外标注的mask千万不要是float类型的0~255范围我在项目里遇到过有人把0和255当作两类去训练结果sigmoid的输出分布跟目标分布完全不匹配训练一直没有意义。4.2 推理流程从模型输出到可视化分割图训练完成后推理阶段要做的事情是加载模型权重把新图像送进网络sigmoid得到概率图然后按阈值转成最终的类别掩码。实际项目里我还会把分割边缘叠加在原图上做可视化这样能直观看到效果。def predict(model, image_tensor, device, threshold0.5): model.eval() with torch.no_grad(): logits model(image_tensor.unsqueeze(0).to(device)) prob torch.sigmoid(logits) mask (prob threshold).float() return mask这里提醒一个细节模型输入图像必须经过跟训练时一模一样的预处理包括尺寸缩放、归一化方式、通道顺序。很多人在推理时忘了归一化直接喂原始像素0~255sigmoid输出的概率分布完全变形分割结果一塌糊涂还以为模型没训练好。还有一个常见问题是灰度图通道数不匹配训练用的3通道推理时给的却是单通道图需要转成RGB再送进去。4.3 可视化训练结果别只盯着Loss曲线看训练过程中除了打印数值指标我还会每5个epoch存一次验证集的分割可视化结果把原图、真实掩码、预测掩码拼成一张三联图看。这个东西的价值比任何指标都直观loss下降不匀速不代表效果差指标不高也不代表分割形状不可用。肉眼观察有几个好处可以看到分割边界是否贴合物体边缘、有没有产生孤立小碎块、预测是不是整体偏移了位置。我从经验里得到的判断标准是如果loss在降但验证Dice一直不涨十有八九是过拟合了如果验证Dice在涨但分割图很碎可能是threshold设得太低试着调到0.6、0.7如果预测整体偏移大概率是数据预处理不一致检查train和val的预处理流程。5. 常见问题与排查技巧实录5.1 训练不收敛的经验排查清单这是我踩过最多的坑列成表格方便你快速定位现象可能原因排查方向Loss完全不动学习率太低 / 输入没归一化打印输入像素范围确认是否在0~1之间Loss剧烈震荡学习率太高 / batch size过小调低lr到1e-4batch size至少8输出全是背景正负样本极度不平衡加入Dice损失或调整权重Validation Dice偏低过拟合加强数据增强或缩小模型通道数训练久不涨BN层在batch size太小时失效换GroupNorm或增大batch size有一种很隐蔽的情况是模型输出logits是float32而mask是int64二者相减时发生隐式类型转换导致数值异常。Pytorch里BCEWithLogitsLoss要求target是float类型CrossEntropyLoss要求target是long类型不同损失函数对标签的要求不一样代码里必须显式转换。我这边的经验是每次喂数据前都强制写mask.float()或mask.long()避免踩类型暗坑。5.2 显存不足的常见应对策略8GB显存跑256x256输入、batch size 16UNet原始参数约3100万训练时会吃紧。我的应对策略本质上就三个字降分辨率。把输入缩到128x128batch size降到8显存占用能下降一大截效果损失不大。如果还想再降把初始通道数从64改成32模型体积直接缩到原来的四分之一。实际上很多公开的分割任务用32通道起步的UNet也完全够用不必执着于原始论文的数字。如果显存还有余量但训练特别慢检查一下DataLoader的num_workers设置。在Linux下我一般设4或8Windows下设0多进程在Windows上容易出问题。pin_memoryTrue也能稍微加速CPU到GPU的数据传输。5.3 上采样尺寸不匹配的边界问题UNet里最经典的一个运行时错误是解码器上采样后的特征尺寸跟编码器的skip特征尺寸差了一两个像素。这个问题的根源是输入图像的尺寸不是2的幂次倍数或者网络层数多导致下采样时奇数尺寸被截断。我的解决方案有两道防线第一道是输入前统一resize到256或512第二道是UpBlock里的pad代码一旦检测到尺寸不一致自动在边缘补零对齐。实测中pad方案确实能保住运行但补零区域会产生一定程度的边缘伪影所以最佳实践仍然是优先保证输入尺寸规范。5.4 数据标注中的坑最后聊几句标注层面的问题。如果你用的是标注工具导出的mask务必确认导出的图是PNG格式且没有经过压缩。我曾经遇到过从标注工具导出jpg格式的mask标签边缘全是压缩伪影形状难看得要命直接导致训练时梯度跨区域乱走。另外mask图里如果存在漏标区域也就是某张图有一部分像素没有对应任何类别这些像素在训练时对损失函数的贡献是纯噪声。稳妥做法是给未知区域单独分配一个类别编号或者在加载数据时把无效像素过滤掉。6. 结尾这个项目后续还能怎么扩展我把这套UNet代码跑通之后做的第一个扩展就是把原始UNet改成带残差连接的版本也就是在DoubleConv的每个卷积之前加一个shortcut收敛速度肉眼可见地加快。再往后我把注意力机制加进去在跳跃连接前接一个注意力门控模块让网络自动学习哪些浅层特征需要被重点保留。这些都是基于现有代码很容易做到的改动如果你上手了基础版完全可以顺着这个方向继续折腾。我个人在实际操作中的体会是不要迷信网上别人的模型代码一定自己一行行敲一遍尤其在forward函数里搞清楚每个张量的形状变化比啥都重要。UNet看起来结构简单但只有亲手写过、跑过、报过错、改对了才算真正掌握。这套代码跑出来的项目X在模拟数据集上的分割效果已经能稳定达到预期验证了从零实现UNet这条路是完全走得通的。最后再分享一个小技巧训练时把每一轮的验证图保存下来按epoch排序看变化你会清晰看到模型从一团模糊慢慢学会勾勒边缘的全过程那个成就感是用任何指标都替代不了的。