U-Net与Attention、R2U-Net等变体对比实验：医学图像分割工程实战

简介面向医学图像分割场景的Unet、AttentionUnet、R2Unet、R2AUet完整可运行代码包配套ISIC 2017皮肤病变数据集适合有一定深度学习基础的算法工程师、研究者和医疗影像分析学习者用于快速复现分割实验、对比不同网络结构在皮肤病变数据上的表现。四种模型分别体现对称编码-解码、注意力门控、残差连接以及残差与注意力结合的改进思路在皮肤病辅助诊断与病灶定位方面具有实际参考价值。包内共211个文件压缩包约25.48MB。其中8个Python脚本为模型核心实现97个png与92个jpg图片构成ISIC数据集的皮肤镜图像及对应掩码另有xml工程配置、sh运行脚本和iml、md说明文件便于快速搭建环境、理解项目结构也可直接查看不同模型在同一批数据上的分割效果。目前已有2159人浏览学习。拿到资源后既能直接运行训练与预测流程也能对照论文分析残差连接、注意力门控等模块对分割精度的影响若想继续深入还可在代码基础上替换数据集或调整超参做进一步算法优化实验。1. 一个开箱即跑的医学图像分割对比实验工程做过分割任务的人都懂最磨人的不是模型本身而是当你想把 U-Net 和它的改进版放在同一套代码里做对比实验时每个版本都得重新配数据、调接口、处理尺寸不匹配的问题。这个资源包把 Unet、AttentionUnet、R2Unet、R2AUet 四种模型集中在同一个工程里附带可直接训练的分割数据集下载解压后按 README 跑通训练就能拿到四组对比结果。适合正在做医学图像分割课题的学生也适合刚接触分割模型、想通过改动代码理解注意力机制和循环残差卷积原理的开发者。2. 工程结构拆解四种模型如何共用一套训练与推理流程2.1 目录组织与核心文件职责解压资源包后工程目录的结构大致如下medical_segmentation/ ├── data/ │ ├── images/ # 原始图像png 格式 │ ├── masks/ # 标签掩码单通道灰度图 │ └── train_val_split.py ├── models/ │ ├── unet.py │ ├── attention_unet.py │ ├── r2unet.py │ ├── r2aunet.py │ └── __init__.py ├── utils/ │ ├── dataset.py # 数据加载与增强 │ ├── losses.py # 混合损失函数 │ ├── metrics.py # Dice 与 IoU 计算 │ └── visualization.py # 预测结果可视化 ├── train.py # 训练入口 ├── predict.py # 单张图像推理 └── config.yaml # 全局配置这类工程设计的核心思路是「模型可替换流程不重写」。在训练脚本里通过字符串映射动态创建模型实例所以切换模型只需要改配置文件中的model_name字段数据加载、损失计算、评估逻辑完全复用。这个设计对做对比实验非常友好——你不需要维护四份彼此独立又高度重复的训练代码。utils/dataset.py中的数据集类继承自 PyTorch 的Dataset在__getitem__里同时读取图像和对应掩码并施加相同的随机变换翻转、旋转、缩放保证增强后图像和标签保持空间对齐。训练集和验证集的划分由train_val_split.py完成按 8:2 比例随机拆分并把划分结果保存为文本文件避免每次训练时划分结果不一致。2.2 训练入口的工作流程train.py是运行时最常打交道的文件。它的流程可以概括为读取配置 → 构建数据集 → 创建模型 → 初始化优化器 → 循环训练 → 保存最优权重。为了适配不同显存大小的显卡config.yaml里暴露了batch_size、input_size、num_workers等关键参数。model_name: r2aunet # 可选 unet / attention_unet / r2unet / r2aunet input_size: 256 # 输入图像尺寸正方形 batch_size: 8 epochs: 120 learning_rate: 0.0001 num_classes: 1 # 二分类分割 loss: bce_dice # 混合损失兼顾像素精度与区域重叠度参数说明input_size决定了训练时所有图像会 resize 到统一尺寸256 是显存占用和性能的折中值追求更高精度可以改为 512但显存需求会明显上升。model_name是切换四种模型的总开关改这一行即可换模型。训练时每轮迭代打印一次 loss、Dice 和 IoU。模型权重的保存逻辑是「验证集 Dice 最高才覆盖保存」训练 120 轮结束后best_model.pth里保存的就是验证集上表现最好的那一轮权重。2.3 数据加载与预处理细节原始数据集的图像尺寸并不统一所以数据加载时需要统一缩放。常见做法是先双线性插值缩放图像对掩码使用最近邻插值——因为掩码是类别标签双线性插值会引入非整数像素值破坏类别语义。from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class SegmentationDataset(Dataset): def __init__(self, image_dir, mask_dir, image_size256, augmentFalse): self.image_paths sorted(image_dir.glob(*.png)) self.mask_paths sorted(mask_dir.glob(*.png)) self.image_size image_size self.augment augment # 图像和掩码文件名一一对应排序保证对齐 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 T.Resize((self.image_size, self.image_size))(image) mask T.Resize((self.image_size, self.image_size), interpolationT.InterpolationMode.NEAREST)(mask) image T.ToTensor()(image) mask torch.as_tensor(np.array(mask), dtypetorch.float32).unsqueeze(0) mask (mask / 255.0).round() # 归一化到 0/1 return image, mask这里的mask / 255.0是二分类任务的关键操作原始掩码像素值为 0 和 255如果不缩放到 0/1 就直接算 BCE 损失模型会学到一个偏置项导致 Dice 指标看起来很怪。将掩码化为 0/1 后模型的输出通过 Sigmoid 激活映射到 (0,1) 区间两者的数值范围一致损失才能正确反映预测质量。3. 四种模型的核心实现从基线到组合改进3.1 经典 U-Net编码器-解码器与跳跃连接U-Net 的骨干由编码器、解码器和跳跃连接三部分组成。编码器通过重复下采样逐步降低特征图分辨率、增加通道数提取语义信息解码器逐步上采样恢复分辨率跳跃连接把编码器各层的细节特征拼接到解码器对应层弥补下采样丢失的空间信息。import torch.nn as nn class DoubleConv(nn.Module): 两个卷积 BN ReLU 的堆叠块 def __init__(self, in_ch, out_ch): super().__init__() self.conv 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.conv(x) class UNet(nn.Module): def __init__(self, in_ch3, out_ch1, features[64, 128, 256, 512]): super().__init__() self.encoders nn.ModuleList() self.pool nn.MaxPool2d(2) self.decoders nn.ModuleList() # 编码器每层一个 DoubleConv层间用 MaxPool 下采样 for f in features: self.encoders.append(DoubleConv(in_ch, f)) in_ch f self.bottleneck DoubleConv(features[-1], features[-1] * 2) # 解码器每层先转置卷积上采样再拼接跳跃连接 for f in reversed(features): self.decoders.append( nn.ConvTranspose2d(f * 2, f, 2, stride2) ) self.decoders.append(DoubleConv(f * 2, f)) self.final nn.Conv2d(features[0], out_ch, 1) def forward(self, x): skip_features [] for enc in self.encoders: x enc(x) skip_features.append(x) x self.pool(x) x self.bottleneck(x) skip_features skip_features[::-1] for i in range(0, len(self.decoders), 2): x self.decoders[i](x) x torch.cat([x, skip_features[i // 2]], dim1) x self.decoders[i 1](x) return self.final(x)参数说明features列表控制每层的通道数默认[64, 128, 256, 512]把列表改为[32, 64, 128, 256]可以生成一个轻量版 U-Net适合显存有限的环境。跳跃连接的实现是torch.cat按通道拼接所以解码器第一个卷积的输入通道是f * 2语义特征和空间细节拼接后统一处理。这个结构的可改性在于池化层换成步长为 2 的卷积可以实现可学习的下采样DoubleConv里的 ReLU 换成 LeakyReLU可以缓解梯度消失——这些改动都不影响整体骨架。3.2 Attention U-Net注意力门控替换朴素拼接Attention U-Net 的改进点很集中不再把跳跃连接直接拼上去而是先经过一个注意力门控Attention Gate对编码器特征做加权筛选。门控信号来自解码器当前层的特征权重越大表示该位置对当前分割任务越重要。class AttentionGate(nn.Module): 注意力门控用解码器特征指导编码器特征加权 def __init__(self, enc_ch, dec_ch, out_ch): super().__init__() self.enc_conv nn.Conv2d(enc_ch, out_ch, 1) self.dec_conv nn.Conv2d(dec_ch, out_ch, 1) self.relu nn.ReLU(inplaceTrue) self.weight nn.Conv2d(out_ch, 1, 1) self.sigmoid nn.Sigmoid() def forward(self, enc_feat, dec_feat): g1 self.enc_conv(enc_feat) # 编码器特征投影 g2 self.dec_conv(dec_feat) # 解码器特征投影 # 逐元素相加并激活再用 1x1 卷积压缩成单通道注意力图 attn self.relu(g1 g2) attn self.weight(attn) attn self.sigmoid(attn) # 归一化到 0~1 return enc_feat * attn # 加权后的特征逻辑说明enc_feat来自跳跃连接dec_feat来自解码器当前层。两者分别经过 1x1 卷积投影到相同通道数后相加再经过 ReLU 和一个 1x1 卷积得到单通道注意力图。Sigmoid 把注意力值限制在 0 到 1 之间最后与enc_feat逐元素相乘相当于对每个位置做了一个软性保留或抑制。在解码器拼接之前插入这个门控的好处非常直观背景区域的编码器特征会被压低权重目标区域的权重接近 1拼接进来的特征信噪比更高。从训练曲线看Attention U-Net 的收敛速度通常比基线 U-Net 快尤其在小目标分割场景下最终 Dice 值更高。3.3 R2U-Net循环残差卷积块替换普通卷积R2U-Net 用「循环残差卷积块」RRC替代了DoubleConv中的普通卷积。循环的意思是同一个卷积块在时间维度上被反复使用多次每次的输出会反馈回输入端再卷积一次形成迭代细化效果残差指块内有一条从输入直接到输出的短接路径保证梯度能顺畅回流。class RecurrentConv(nn.Module): 单次卷积 BN ReLU供循环块反复调用 def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class RecurrentResidualBlock(nn.Module): def __init__(self, in_ch, out_ch, t2): super().__init__() self.t t self.conv1 RecurrentConv(in_ch, out_ch) self.conv2 RecurrentConv(out_ch, out_ch) def forward(self, x): out self.conv1(x) for _ in range(self.t - 1): out self.conv1(out) residual out out self.conv2(residual) for _ in range(self.t - 1): out self.conv2(out) return out x if x.shape out.shape else out self.shortcut(x)逻辑说明循环次数t是 R2U-Net 的超参数默认取 2表示同一个卷积层被调用两次。conv1的输入是原始特征x第一次输出后第二次再把结果喂回同一个conv1——权重是共享的这相当于让卷积核在相同感受野内多一次特征提炼。残差连接把块输入端直接加到输出端缓解循环加深带来的梯度衰减。我这里省略了shortcut层的定义实际工程中如果输入输出通道不一致需要补一个 1x1 卷积对齐维度。R2U-Net 的代价是训练速度变慢——循环次数让前向计算量翻倍但精度收益在数据集较小或标注噪声较大时很明显因为循环结构相当于隐式数据增强。3.4 R2A-U-Net循环残差与注意力的组合R2A-U-Net 的做法不是把两种改进简单堆叠而是让注意力门控作用于循环残差块的输出。编码器侧用 RRC 提取特征解码器侧同样用 RRC 恢复分辨率跳跃连接处插入 Attention Gate。class R2AU_Block(nn.Module): def __init__(self, in_ch, out_ch, t2): super().__init__() self.rrc RecurrentResidualBlock(in_ch, out_ch, t) self.attn AttentionGate(out_ch, out_ch, out_ch) def forward(self, enc_feat, dec_feat): enc_feat self.rrc(enc_feat) gated self.attn(enc_feat, dec_feat) return gated组合的动机是RRC 负责把浅层的纹理信息迭代细化Attention Gate 负责筛选哪些细化后的特征值得传给解码器。如果只加注意力不加循环小目标边缘的细节特征会在早期层丢失只加循环不加注意力背景噪声会被同等放大。两者组合后模型既有了深层次的特征提炼能力也有了抑制无关区域的选择能力。从实验角度看R2A-U-Net 的参数量是四种模型里最大的训练耗时也最长。如果你的显存有限可以把features改为[32, 64, 128, 256]或者把循环次数t降为 1——循环 t1 时 RRC 退化为普通残差卷积R2A-U-Net 就变成了「残差 注意力」的轻量组合仍然比纯 U-Net 效果好。4. 训练与推理实战从数据集到分割结果4.1 数据集准备与划分检查资源包附带的数据集已经按images和masks分好目录文件名一一对应。第一次跑训练前建议先运行数据划分脚本并检查输出python data/train_val_split.py --image_dir data/images --mask_dir data/masks --val_ratio 0.2脚本会在data目录下生成train.txt和val.txt每行是一个不带扩展名的文件名。手动打开这两个文件随机挑几行去images目录确认对应文件存在。这一步虽然简单但能避免训练跑起来后因为文件名不匹配读到空数据。4.2 训练启动与损失函数解析训练命令如下python train.py --config config.yaml默认配置下120 轮训练在单张 1080Ti 级别显卡上大约需要 23 小时。资源包使用的损失函数是 BCE 与 Dice 的加权组合class BCEDiceLoss(nn.Module): def __init__(self, bce_weight0.5, dice_weight0.5): super().__init__() self.bce_weight bce_weight self.dice_weight dice_weight def forward(self, pred, target): bce nn.functional.binary_cross_entropy_with_logits(pred, target) pred_prob torch.sigmoid(pred) smooth 1e-5 intersection (pred_prob * target).sum() dice 1 - (2 * intersection smooth) / (pred_prob.sum() target.sum() smooth) return self.bce_weight * bce self.dice_weight * dice参数说明bce_weight和dice_weight控制两项损失的占比。BCE 逐像素计算损失对类别不平衡敏感Dice 关注预测区域和真实区域的重叠度。默认各占 0.5 在大多数分割任务上表现稳定。如果你的目标区域在图像中占比很小可以把dice_weight提高到 0.7。训练日志里重点关注验证集 Dice 的走势。正常情况是前 20 轮快速上升40 轮后进入缓慢提升阶段此时可以停止训练不必等满 120 轮。4.3 单张图像推理与可视化训练完成后用predict.py对单张图像做推理python predict.py --image data/images/sample.png --weight best_model.pth --model r2aunet推理脚本内部的处理流程是读图 → 缩放到 256x256 → 转张量 → 模型前向 → Sigmoid → 阈值 0.5 转二值掩码 → 原图与掩码叠加保存。可视化脚本会把三张图并排输出到output/目录原图、预测掩码、原图上叠加红色半透明掩码的融合图。判断结果好坏时除了看 Dice 数值一定要肉眼检查融合图。常见问题是预测掩码边缘比真实掩码粗一圈这是上采样恢复分辨率时边界信息丢失造成的如果预测结果出现大量孤立小点通常是训练轮数不够或数据集噪声大把num_classes从 1 改为 2、采用 WeightedCrossEntropyLoss 能缓解。5. 避坑指南跑通这个工程的常见问题排查5.1 显存溢出报错 OutOfMemoryError现象训练跑到第 3 个 epoch 时程序报错退出提示CUDA out of memory。原因默认input_size256、batch_size8在 6GB 显存显卡上跑 R2A-U-Net 会超限。四种模型里 R2A-U-Net 的显存占用最高因为它同时具备循环卷积和注意力分支。解决把config.yaml中的batch_size改为 2 或 4如果还不够就把input_size改为 192。改input_size后数据集会自动重新缩放不需要手动预处理。一个折中方案是用梯度累积——每 4 步反向一次效果等价于 batch_size32但显存占用不变。5.2 训练时 loss 正常但预测结果全黑或全白现象训练曲线平滑下降验证集 Dice 也显示在提升但把预测图保存下来发现整张图是一个颜色没有任何分割轮廓。原因极可能数据归一化出了问题。掩码归一化到 0/1 后模型输出经过 Sigmoid 后默认阈值 0.5 判断前景背景。如果原始图像没有做标准化输入像素值范围在 01而模型权重初始化基于标准正态分布前几层输出的特征分布会偏移导致 Sigmoid 输出永远接近 0 或接近 1。解决在dataset.py中为图像增加标准化步骤使用均值[0.485, 0.456, 0.406]和方差[0.229, 0.224, 0.225]ImageNet 统计值。加完标准化后重新训练预测结果会正常。另外确认掩码保存格式是否为 pngjpeg 压缩会在掩码边缘引入灰色像素阈值化后产生锯齿。5.3 训练前期 Dice 损失出现 NaN现象第 1 个 epoch 的 loss 输出是nan之后一直维持nan。原因最常见的是优化器学习率过大导致梯度爆炸。另一个可能性是损失函数的平滑项太小——smooth1e-5在目标区域极小的情况下分母接近 0浮点溢出。解决先把learning_rate从0.0001降到0.00001确认训练能启动后再逐步调大。对于 Dice 损失把smooth从1e-5改为1.0能有效避免除零问题而且几乎不影响最终指标。可以在训练脚本里添加一个梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)从机制上防止梯度过大。5.4 训练集 Dice 很高但验证集 Dice 波动剧烈现象训练集 Dice 在第 60 轮达到 0.95验证集最好成绩只有 0.82且每轮之间波动 23 个百分点。原因这不是代码 bug而是数据集划分和增强策略导致的过拟合。如果验证图像和训练图像来自同一批次拍摄、背景高度相似验证集 Dice 会虚高反之如果拍摄条件差异大泛化性能就体现在波动上。解决在train/val_split.py中按「图像采集时间」分桶划分而不是随机划分保证同一时间段的图像不出现在两个集合中。另外把数据增强中的旋转角度从 90 度改为 ±15 度旋转过大会让模型学到错误的几何不变性反而影响真实场景效果。5.5 更换数据集后训练报错通道数不匹配现象把自己收集的数据集放入data/images和data/masks后训练报错size mismatch模型参数加载失败或前向计算维度对不上。原因资源包默认输入是三通道 RGB 图像in_ch3和二分类输出out_ch1。如果你的图像是单通道灰度图如某些医学影像Conv2d(3, 64)会直接报错如果你的掩码是多类别标注out_ch1会让多类信息丢失。解决单通道灰度图在dataset.py里改为image.convert(L)并顺带把模型初始化参数改为UNet(in_ch1)。多类别分割要把out_ch改为类别数损失函数从 BCE 换为 CrossEntropyLoss同时冻结最后一层卷积的预训练权重——常见做法是加载预训练权重后跳过final层的state_dict。6. 二次开发思路替换数据集、改造模型与增强评估这套工程最大的价值不在于开箱即跑而在于替换成本低。我一般在拿到这类项目后第一件事就是在不改内部逻辑的前提下把数据集换成自己的数据验证数据管道的正确性再谈改模型。换数据集时只需要保证两个条件图像和掩码文件名严格一致掩码是单通道 png 格式。如果数据是 nii.gz 格式的医学影像CT/MRI 分割场景在dataset.py的__getitem__里用 SimpleITK 读取并切片转成 numpy 数组即可不需要改动模型。import itk import numpy as np def load_nii_as_image(path, slice_idx): img itk.imread(str(path)) arr itk.array_from_image(img) # 取中间切片转成 HWC 格式 slice_2d arr[slice_idx] return slice_2d改模型方面常用的路径是给编码器换上预训练的 ResNet 骨干。做法是把models/unet.py中自定义的encoders替换为torchvision.models.resnet34(pretrainedTrue)的layer0到layer3把跳跃连接的通道数按 ResNet 的输出通道数设置。这个改动约 30 行代码但能明显提升分割精度因为 ImageNet 预训练特征对自然图像有很强的迁移能力。我一般还会加上 mIoU 评估指标补全对比实验的说服力def calculate_miou(pred, target, num_classes): ious [] for cls in range(num_classes): pred_mask (pred cls) target_mask (target cls) intersection (pred_mask target_mask).sum() union (pred_mask | target_mask).sum() if union 0: continue ious.append(intersection / union) return np.mean(ious)这段代码按类别逐个计算交并比。Dice 只反映前景区域的重叠度而在多类别分割中背景类占比大Dice 会被背景虚高。mIoU 把所有类别平等对待能暴露出模型对少数类别的识别短板——这是我做对比实验时最依赖的一个指标。四类模型跑完一遍后我的复盘习惯是先看验证集 mIoU 排名再看预测可视化的边缘质量。U-Net 基线往往排名垫底但训练最快R2A-U-Net 指标最好但显存和训练成本翻倍。真正的工程建议是先跑 U-Net 确认数据管道没问题再用 R2A-U-Net 追求精度上限如果上线部署有资源限制再退到 Attention U-Net 做精度与速度的平衡。从那以后我每次复现这类工程都强制自己先画一条基线跑通全流程再逐层叠加改进避免一上来就追 SOTA 结果结果翻车在数据预处理上。希望帮到你。本文还有配套的精品资源点击获取