U-Net家族四模型对比：从Attention到R2AUet的医学图像分割实践

简介Unet、AttentionUnet、R2Unet、R2AUet四种图像分割模型的完整可运行实现配套ISIC 2017皮肤病变数据集面向计算机视觉与医学影像分析领域的学习者、研究者和算法工程师适合从原理理解到实验复现的进阶练习。代码以Python文件为主涵盖网络结构、训练脚本与配置数据集由皮肤镜图像构成jpg与png格式的图片占绝大多数另有xml标注、md说明文档及sh脚本辅助使用方便直接搭建训练环境。压缩包共211个文件整体大小25.48MB目录结构清晰便于按模型或数据模块检索。目前已有2159人学习下载。通过对比残差块与注意力门控在Unet骨架中的融合方式读者能够理清四种模型的演进脉络掌握医学图像分割中的预处理、损失函数、优化器选择与超参数调优思路并可在ISIC 2017数据上直接训练、验证与评估为论文复现或实际项目落地提供可参考的代码基线。1. 从四个网络一起调通说起这个项目解决的是图像分割选型难题拿到这份源码包的那一刻我其实没抱太高期望因为标题里的 Unet、AttentionUnet、R2Unet、R2AUet 四个模型看起来就像是同一套东西换了几层皮。可真在本地把训练跑起来后发现它们的分割行为差异非常明显普通 U-Net 边缘糊、Attention U-Net 能把病灶边界抠得更干净、R2U-Net 对小目标更敏感而 R2AUet 在两者之间做了折中。这个项目把四种模型、训练/验证/预测脚本和一份可直接加载的数据集都打包在了一起装好依赖后不用改代码就能跑出分割结果。对做医学图像分割、遥感地物识别或工业缺陷检测的开发者来说它最大的价值是省掉了找数据、对齐接口、调试模型结构的大量重复劳动。2. 工程结构拆解模型、数据、训练三者如何分层2.1 目录结构与职责划分先看整体布局。这份资源的目录结构非常常见几乎就是标准 PyTorch 工程的分层方式方便直接迁移到别的项目里。我用 tree 命令扫了一遍骨架长这样unet-family/ ├── models/ │ ├── unet.py │ ├── attention_unet.py │ ├── r2unet.py │ └── r2attention_unet.py ├── dataset/ │ ├── data_loader.py │ └── data/ ├── utils/ │ ├── losses.py │ ├── metrics.py │ ├── visualize.py │ └── logger.py ├── config.py ├── train.py ├── predict.py └── requirements.txtmodels 目录下每个文件只放一个模型类类的对外接口统一为__init__(in_channels, out_classes, base_channel)这样在 train.py 里可以用字典映射动态实例化不需要在每个脚本里写死模型类名。dataset/data_loader.py 里实现的是继承torch.utils.data.Dataset的加载器返回的是(image, mask)元组mask 是单通道灰度图值为 0 或 1背景和前景没有混淆。utils 里把损失函数、Dice 系数计算和图像保存拆开训练日志也只打印到控制台想接 TensorBoard 可以再自己加。需要留意的点是config.py它是整个工程的全局配置中心。训练参数和路径都集中在这里而不是散落在各脚本里。我一般拿到项目后先改这个文件而不是去翻训练脚本。2.2 技术选型为什么用 PyTorch依赖文件 requirements.txt 里写的是torch 1.8、numpy、opencv-python、Pillow、tqdm都是图像分割生态里最常用的库。用 PyTorch 的一个好处是动态图调试方便特别是当你想改模型内部某个层的输出尺寸时只要在 forward 里加一行打印就能立刻看到张量形状。另一好处是预训练权重容易加载比如把 ResNet 当作 U-Net 的编码器主干时直接用torchvision.models的权重就行这个项目里没有这么做但如果你二次开发想升级路径是完全畅通的。数据集如果没有特殊标注用 opencv 读取图像和掩膜就足够了不需要引入 SimpleITK 这类重型医学库如果换成医学影像原格式比如.nii那时再增加相应依赖。2.3 数据加载与预处理管线数据集部分不是玩具它包含了已经划分好的训练集和测试集文件列表。data_loader.py 核心流程分为三步先从config.py里读取图像文件夹路径再对图像做三通道归一化对 mask 做二值化最后执行在线数据增强。增强只使用随机翻转、随机旋转和缩放没有用复杂的弹性畸变原因是为了保证 mask 和图像同步变化减少标注错位风险。关键代码如下class SegDataset(Dataset): def __init__(self, img_list, mask_list, transformNone, image_size(256, 256)): self.img_list img_list self.mask_list mask_list self.transform transform self.image_size image_size def __len__(self): return len(self.img_list) def __getitem__(self, idx): image cv2.imread(self.img_list[idx]) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image cv2.resize(image, self.image_size) image (image.astype(np.float32) / 255.0) # 归一化到 0~1 mask cv2.imread(self.mask_list[idx], cv2.IMREAD_GRAYSCALE) mask cv2.resize(mask, self.image_size, interpolationcv2.INTER_NEAREST) mask (mask 127).astype(np.float32) # 二值化保留前景 if self.transform is not None: image self.transform(image) # CHW return torch.from_numpy(image.transpose(2, 0, 1)), torch.from_numpy(mask)这里有两个很容易被忽略的参数image_size必须和模型里下采样后的特征图尺寸匹配否则后面的跳跃连接拼接会报错而INTER_NEAREST在缩放 mask 时是必须的如果用线性插值mask 会出现中间灰度值导致 loss 在取值 0.5 的地方震荡。很多新手在这两个细节上翻车训练时报 shape 不匹配或者 Dice 指标奇怪地卡在 0.5。3. 核心模型拆解从 U-Net 到 R2AUet 的四级递进3.1 U-Net 的基本骨架U-Net 的代码是所有变体的地基先把它看透后面就好理解。这里的实现用了一个DoubleConv作为基本单元每个单元包含两次卷积、BN 和 ReLU。编码器每下采样一次通道数翻倍特征图尺寸减半解码器每次上采样后把编码器对应层的特征图沿通道拼接起来形成 U 形结构。核心 forward 如下def forward(self, x): features [] # encoder for block, pool in zip(self.down_blocks, self.pools): x block(x) features.append(x) x pool(x) x self.bottleneck(x) # decoder for up, block, skip in zip(self.up_blocks, self.down_blocks[::-1], features[::-1]): x up(x) x torch.cat([x, skip], dim1) # 跳跃连接 x block(x) return x这里跳跃连接用的是通道拼接而不是加法这是 U-Net 最核心的设计。拼接后的通道数是解码器上采样结果和编码器特征图通道数之和所以DoubleConv的第一个卷积需要动态接收这个拼接后的通道数。在这个实现里编码器各层输出通道分别为 64、128、256、512上采样后分别是 64、128、256、512拼接后对应 128、256、512、1024这就是各解码块第一层卷积输入维度。如果数据集是 3 通道输入第一层卷积输入就是 3如果换成灰度单通道只需要把in_channels改成 1。3.2 Attention U-Net门控信号自动聚焦Attention U-Net 最大的变化是在跳跃连接前增加了一个注意力门Attention Gate。它不直接拼接编码器特征而是先用解码器的上采样特征作为门控信号计算一个 0~1 的注意力权重然后乘到编码器特征上。这样做相当于告诉网络跳跃连接里哪些区域是对当前目标有用的哪些是背景噪声。代码实现通常是把编码器特征gating_signal和解码器特征skip分别做 1x1 卷积相加后过 ReLU再过 Sigmoid最后与skip相乘。class AttentionGate(nn.Module): def __init__(self, in_channels, gate_channels, out_channels): super().__init__() self.W_g nn.Conv2d(gate_channels, out_channels, kernel_size1) self.W_x nn.Conv2d(in_channels, out_channels, kernel_size1) self.psi nn.Conv2d(out_channels, 1, kernel_size1) def forward(self, x, g): g_resized F.interpolate(g, sizex.shape[2:], modebilinear, align_cornersTrue) theta_x self.W_x(x) theta_g self.W_g(g_resized) f F.relu(theta_x theta_g) sig torch.sigmoid(self.psi(f)) return x * sig在这个资源里注意力门的输入是解码器当前层上采样后的特征输出是经过空间加权后的编码器特征。interpolate这一步是必须的因为解码器特征尺寸和编码器特征尺寸往往差一倍需要先把门控信号缩放到与跳跃连接相同的尺寸。这个模块引入后参数量增加不大但分割边界会明显更平滑因为模型学会了忽略背景区域。3.3 R2U-Net循环残差卷积R2U-Net 用循环卷积模块Recurrent Convolutional Block替换了普通 DoubleConv。它把同一个卷积层重复执行 T 次每次的输入包含上一次的输出和原始输入最后再加上一个残差连接。这样做的好处是在不增加太多参数的情况下加深了感受野对小目标和细长结构更友好。核心替代块是R2Convclass R2Conv(nn.Module): def __init__(self, in_ch, out_ch, T2): super().__init__() self.conv nn.Conv2d(in_ch, out_ch, kernel_size3, padding1) self.bn nn.BatchNorm2d(out_ch) self.relu nn.ReLU(inplaceTrue) self.T T self.residual nn.Conv2d(in_ch, out_ch, kernel_size1) if in_ch ! out_ch else nn.Identity() def forward(self, x): identity self.residual(x) out self.relu(self.bn(self.conv(x))) for _ in range(self.T - 1): out self.relu(self.bn(self.conv(out))) return out identityT就是循环次数资源里默认设置成 2。很多人以为 T 越多越好实际上T4时参数量并不线性增长但深度增加会导致梯度消散风险变大。如果你在自己数据集上训练发现训练曲线波动可以先把 T 调回 2 试试而不是改学习率。3.4 R2AUet组合变体R2AUet 是资源里最完整的模型将循环残差卷积和注意力门两种机制合在一起。编码器、瓶颈、解码器用的都是 R2Conv跳跃连接处挂上 AttentionGate。这样做的效果是既聚合了多尺度上下文又对无关区域做抑制。这个模型的参数量是四个里最大的训练时显存占用也会增加但效果通常最好。它在扩充分支上使用x * sig加权这个sig值是学习出来的不会出现硬阈值。因此你的数据如果前景和背景像素比例严重失衡R2AUet 往往能给出更稳的结果因为它能自动把背景区域权重压低。4. 训练与评估一条命令跑通到拿到分割掩膜4.1 从配置文件到启动训练修改好config.py里的数据路径然后执行python train.py --model r2unet --epochs 60 --batch-size 4训练就开始了。train.py 的开头部分会先打印模型名称、参数量、数据集样本数这些信息对定位问题非常关键。模型的选择通过一个字典映射MODEL_MAP { unet: UNet, attention_unet: AttentionUNet, r2unet: R2UNet, r2auet: R2AttentionUNet, } model MODEL_MAP[args.model](in_channels3, out_classes1, base_channel64)base_channel64决定了第一个卷积层的输出通道数后面每层的通道数都是它的倍数。如果你的显卡显存不够把base_channel降到 32参数量和显存占用会迅速减半。训练时每过一个 epoch程序会保存当前 epoch 的权重到checkpoints/目录命名为best_dice.pth保存标准是验证集 Dice 不再提升。4.2 损失函数与评估指标的计算方式损失函数用的是BCEWithLogitsLoss加 Dice Loss 的线性组合权重各 0.5。只用 BCE 时由于背景像素远多于前景模型会倾向把全部预测为背景Dice 指标显得低。加入 Dice Loss 后模型被迫关注前景区域具体实现如下class HybridLoss(nn.Module): def __init__(self): super().__init__() self.bce nn.BCEWithLogitsLoss() def forward(self, logits, targets): bce self.bce(logits, targets) probs torch.sigmoid(logits) flat_probs probs.view(probs.size(0), -1) flat_targs targets.view(targets.size(0), -1) dice 1 - (2 * (flat_probs * flat_targs).sum(1) 1e-6) / \ (flat_probs.sum(1) flat_targs.sum(1) 1e-6) return bce dice.mean()这里1e-6是平滑项用来防止分母为零。BCE 和 Dice 的尺度不同通常 BCE 在 0.2~0.7Dice 在 0~1直接相加时 Dice 权重会偏大所以实际使用也可以给 Dice 加一个系数 0.5。训练脚本里已经调试好了权重你不一定需要改只有当损失值出现异常波动时才回头调整。评估指标除了 loss 外脚本还计算了验证集的 Dice 和 IoU。IoU 直接从预测概率和真实 mask 计算预测概率大于 0.5 视为前景然后统计交集和并集。4.3 推理与可视化predict.py 可以把测试集图片批量处理输出结果到output/。它会将原始图像、预测概率图、二值掩膜、原图叠加掩膜四张图拼成一张大图保存方便人工检查。如果你想单独导出掩膜修改save_mask_only True即可。推理时注意要保证输入尺寸和训练时一致否则测试结果会变形。如果你在 predict.py 里看到model.eval()配合torch.no_grad()这是标准做法可以避免 Dropout 和 BN 在推理时造成输出振荡。5. 避坑与常见问题把可运行项目改到自己环境中的实战排查5.1 跳跃连接时 shape 不匹配现象运行训练脚本时出现RuntimeError: Sizes of tensors must match指向torch.cat一行。原因输入图片尺寸不是 2 的整数次幂倍数。U-Net 编码器有 4 次下采样每次下采样让宽高减半如果图片尺寸是 257x257最后一次下采样会出现尺寸除不尽导致上下采样后的特征图宽高对不上。数据集目录里的原始图片有不少是任意尺寸data_loader 里虽然做了 resize但若标定的image_size和模型步长不匹配就会触发这个问题。解决把config.py中的image_size设为 256 或 512 这样的 2 的幂次。不要使用paddingsame去规避那会改变整个下采样步长让后面的注意力门尺寸也乱掉。我习惯在模型 forward 里打印一下每层尺寸快速定位是哪一层出的问题。5.2 显存溢出batch size 必须调小现象执行python train.py --model r2auet --batch-size 8后几秒内报CUDA out of memory。原因R2AUet 的循环卷积会展开多个时间步中间激活值比普通 U-Net 多许多8 张 256x256 图同时过模型显存占用大约 6GB如果显卡是 4GB 版本必然溢出。解决先降 batch size 到 2再降base_channel到 32。两步都做完显存占用能压到 2GB 以内。如果还需要更大批尺寸启用梯度累积每 4 步更新一次参数效果接近 batch size8。注意torch.cuda.empty_cache()在每次验证后调用否则显存碎片会累积。5.3 数据集加载失败图片路径里带有中文或相对路径问题现象FileNotFoundError: No such file or directory并且文件名是中文。原因项目默认环境是 Linux 下的 UTF-8但 Windows 下部分 OpenCV 版本读取中文路径会失败这是 cv2.imread 的已知表现。解决把整个项目目录放到纯英文路径下数据集内部文件夹也全部改成英文名。如果你不想动目录可以在 data_loader 里用np.fromfile加cv2.imdecode替代cv2.imread我一般会同时保留两种读取逻辑这样别人拿到项目在不同系统上跑都不会报错。5.4 训练不收敛Dice 指标一直横盘现象loss 下降很快头 5 个 epoch 就已很低但验证集 Dice 一直停在 0.6 附近。原因大概率是 mask 的归一化出了问题。有的公开数据集 mask 里背景像素是 255前景是 0而这份资源里的样例数据希望背景为 0、前景为 1如果直接读进去不取反模型学到的前景其实是背景Dice 却会暂时维持在一个不高不低的值。解决在__getitem__之后打印一行mask.unique()看 mask 是否只有 0 和 1 两个数值。如果出现255先做mask (mask 255)再转浮点。这种错误非常隐蔽因为 loss 在降但你实际要分割的东西完全没被学会。5.5 加载权重时 strict 报错现象换显卡继续训练时model.load_state_dict(torch.load(best_dice.pth))报Missing key(s) in state_dict。原因上一次训练的base_channel或分类数out_classes和当前模型不一致。这是最常见的历史遗留问题不是代码 bug。解决加载权重前打印state_dict里第一个卷积层的 shape和当前模型对比。如果不想改全部结构可以把load_state_dict(..., strictFalse)但这会静默丢弃不匹配的层可能导致指标神秘下降。正确做法是回到config.py把base_channel、out_classes改成和之前相同的值然后重新加载。6. 二次开发换数据集、加自定义模块与导出部署6.1 把自己的数据集接进来换数据集的成本不大主要是把数据整理成训练脚本认识的格式。如果你有一堆原图和一个 CSV 文件标明每张图的 mask 路径直接改写dataset/data_loader.py里的初始化函数即可。我习惯在__init__里加一个参数csv_path然后读取成两个列表。注意 mask 不需要在 Dataset 里做 One-Hot因为当前是二分类如果你想做多分类比如背景、肝脏、肿瘤三类就把mask的像素值改成 0/1/2并把out_classes设为 3最后在HybridLoss里把 BCE 换成交叉熵Dice 改成对每个类别分别计算再取平均。这个改动大约半天能完成。# 读取自定义标注的示意 def _parse_csv(self, csv_path): df pd.read_csv(csv_path) self.img_list df[image].tolist() self.mask_list df[mask].tolist()接好之后第一次跑通先用 20 张图像验证流程不要直接全量训练。确认 mask 的像素值、图像通道顺序、resize 尺寸、预测输出的值域都符合预期再放大训练规模。6.2 在跳跃连接上加一个额外的深度监督原模型的输出只有一个。你可以给解码器的每层都加一个 1x1 卷积输出预测图然后和真实 mask 计算损失再把多个损失相加。这种方式叫深度监督特别适合改善深层网络的梯度流动。实现也不复杂在 forward 中收集中间输出predict 时只保留最后一层。def forward(self, x): ... outputs [] for up, block, skip in zip(...): x up(x) x block(torch.cat([x, skip], dim1)) outputs.append(self.heads[len(outputs)](x)) # 每个解码层一个轻量head return outputs[-1], torch.stack(outputs) # 主输出 深度监督列表参数上值得注意辅助损失权重不要设太大0.2 比较合适太重会让浅层卷积过于专注边界细节而忽略语义。训练时你可以在HybridLoss里把辅助损失乘上 0.2再和主损失相加。这个技巧在目标很小、背景很大的分割任务上非常有用R2AUet 本身已经具备较强的梯度回传能力但这个改动仍能继续提升 13 个点的 Dice。6.3 导出 ONNX 并做半精度推理模型调好之后如果要做推理加速我会先把 PyTorch 权重导出成 ONNX。这里有一个注意点模型里的F.interpolate在导出时需要指定opset_version11否则某些算子可能不被支持。导出脚本的核心只有几行dummy_input torch.randn(1, 3, 256, 256) model.cpu().eval() torch.onnx.export( model, dummy_input, model.onnx, input_names[image], output_names[mask], dynamic_axes{image: {0: batch}, mask: {0: batch}}, opset_version11 )导出后可以用 ONNX Runtime 加载推理速度通常比原生 PyTorch 快 30%50%。如果想继续压榨性能把输入转成半精度fp16要注意 BN 层在半精度下容易出现数值不稳定建议把 BN 的 eps 调大一点点比如从1e-5改成1e-4。这样导出后精度基本不下降显存占用还能再降一半。实际拆这个项目时我踩过的坑大多集中在数据读取和尺寸对齐上模型本身的代码写得比较规范没有太多要返工的地方。现在每接到一个新的分割任务我都会强制先跑一遍predict.py的基线结果再决定要不要上注意力门或者循环残差。调参再久不如先把四种模型各自的底线能力摸清楚。这个项目把实验成本降到了最低希望帮到你。本文还有配套的精品资源点击获取