多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

U-Net/Attention U-Net/R2U-Net/R2AUet一次跑通:分割网络PyTorch实现与对比

U-Net/Attention U-Net/R2U-Net/R2AUet一次跑通:分割网络PyTorch实现与对比 简介面向图像分割学习与医学影像研究的完整可运行代码包整合Unet、AttentionUnet、R2Unet、R2AUet四种经典网络并配套ISIC 2017皮肤病变公开数据集覆盖皮肤病图像的像素级分割与多个模型的效果对比。适合计算机视觉初学者入门学习也适合医疗影像开发者直接训练、调参和复现实验。压缩包共211个文件以97张PNG和92张JPG皮肤镜图像为数据主体另有8个Python模型与训练脚本、5个XML配置、5个编译后的pyc文件及工程说明文档整体约25MB轻量便于快速下载部署。目前平台已有2155人学习下载。包内目录结构清晰包含训练脚本、预处理逻辑与可直接运行的工程通过对比实验可深入理解注意力机制与残差连接对分割精度的实际提升也能直接支撑课程设计、论文实验或辅助诊断研究并为后续改进提供完整基线。1. 四个分割网络一次跑通为什么这个组合是分割实验最实用的基线起步如果你的项目标题里同时躺着 U-Net、Attention U-Net、R2U-Net、R2AUet并且还带上了“可运行、包含数据集”这七个字那说明你要的不是一篇原理科普而是一套能立刻跑起来、能出对比图的代码底座。这个组合在医学图像分割、遥感语义分割、道路裂缝检测里出现频率极高原因很简单它们共享同一个编码器-解码器骨架差异只发生在跳跃连接和卷积块内部因此可以在同一份代码里通过开关切换非常适合用来做 baseline 对比或为自己提出的改进模型垫底。我一开始接触这套组合是为了给息肉分割数据集跑一组对照实验当时最大的困扰不是理解网络而是找不到一份能同时输出四个模型指标、又不用各自维护四套不同数据加载逻辑的代码。后来我把它们收敛到了一个工程框架里才发现这个组合真正的价值不在某个网络精度多高而在于你能在完全相同的训练条件下回答一个问题在同样的数据和超参下注意力机制和循环残差卷积到底各自贡献了多少涨点。这篇就按我实际搭建的顺序来讲从结构差异、代码骨架、训练踩坑到验证技巧尽量让你照着能一次跑通。2. 四兄弟的结构差异从跳跃连接到注意力门改动到底在哪2.1 U-Net所有变体的骨架跳跃连接是全部的胜负手原始 U-Net 的思路可以拆成三句话编码器逐层下采样提取语义特征解码器逐层上采样恢复空间分辨率跳跃连接把同尺度的编码器特征拼到解码器特征上。这个设计在当时解决了一个很实际的问题——图像分割要求逐像素输出但连续池化会丢失边界细节跳跃连接等于是给解码器开了一扇直达浅层特征的“后悔药窗口”。实现层面U-Net 的每个层级通常是一个双卷积块两个 3x3 卷积 ReLU BatchNorm下采样用 stride2 的卷积或 maxpool上采样用转置卷积或双线性插值。拼接方式用的是通道维度 concat而不是逐元素相加这也是它和后来 ResNet 风格分割网络最明显的区别——concat 保留了编码器特征的全部信息但会带来更大的通道数计算和显存开销也随之上升。我一般会把 U-Net 当作整套实验的“标尺”网络所有其他变体都和它保持相同的编码器通道配置、相同的上采样方式只做局部替换。这样最后画对比表时精度差异才能归因到结构改动本身而不是因为某个网络用了更强的骨干而占了便宜。做研究对比时这一点尤其重要否则审稿人或导师第一个问题就是“你的控制变量呢”。2.2 Attention U-Net注意力门在跳跃连接上做“软过滤”Attention U-Net 对 U-Net 的改动只有一个点在跳跃连接送入 concat 之前加了一个注意力门Attention Gate。这个门的输入有两路——来自编码器的浅层特征和来自解码器的深层特征通过一个 1x1 卷积加 sigmoid 生成一张空间注意力权重图再用这张权重图对浅层特征逐元素加权。换句话说模型学会了在融合时主动忽略背景区域把注意力集中到目标区域上。这个改动对医学图像分割特别友好因为这类数据的背景占比往往很大目标可能只占整张图的百分之几甚至千分之几。我在息肉分割数据上观察到的现象是U-Net 在背景区域会产生零碎的假阳性而 Attention U-Net 的假阳性位置会明显向真实目标边缘收敛。代价也不是没有注意力门额外引入了可学习参数和一次 sigmoid 计算训练前期收敛会略慢因为门控参数需要先“学会”哪里该注意。注意力门的 PyTorch 实现比较直接它内部就是两个 1x1 卷积、一次加法、一个 ReLU、一个 sigmoid 和一个重采样操作。后面章节我会给完整代码块。这里先记住一个关键点注意力门的输入特征是编码器特征和解码器特征两者分辨率必须一致所以实际使用时需要在跳跃连接前对深层特征做上采样或者让门内部做插值对齐。2.3 R2U-Net把普通卷积换成循环残差卷积RRCNN 块R2U-Net 的思路来自另一个方向它认为 U-Net 的卷积块提取特征不够充分于是把双卷积块替换成了循环残差卷积块Recurrent Residual Convolutional Neural Network block简称 RRCNN。这个块内部做了一件很聪明的事同一个卷积层递归地作用多次每多递归一次等效感受野就往外扩一圈同时保留了残差连接让梯度能直接回流到块的输入端。我第一次看到这个结构时觉得它有点反直觉——循环共享权重意味着参数量不变但计算量翻倍。实际效果确实值得这个代价R2U-Net 在血管分割、细胞分割这类纹理细节密集的任务上往往会比 U-Net 更稳定地不丢小目标。原因也很好理解循环卷积相当于在每个位置做了多次特征再提炼相当于每个像素都见过领域内更大范围的信息。它的实现难点在于循环次数通常设为 2 或 3和内部结构一个 RRCNN 块内部包含两层循环卷积和一个残差连接。循环次数是超参不是越大越好我试过从 2 调到 3训练时间长了约三成精度只涨了不到一个点所以默认设 2 就够了。2.4 R2AUet循环残差 注意力门感受野与抑制噪声同时发力R2AUet 是 R2U-Net 和 Attention U-Net 的简单叠加编码器里的普通卷积块换成 RRCNN 块跳跃连接上再加注意力门。这个名字里没有多余的设计它就是把前面两个改进塞进了同一个网络——用循环残差卷积增强特征提取能力用注意力门抑制无关区域的响应。从实验角度看R2AUet 几乎总是四个网络里表现最好的一个但单次实验的涨幅可能不大。我跑过两组对比一组是息肉分割R2AUet 比 U-Net 的 mIoU 高了约两个百分点另一组是道路裂缝分割涨点甚至不到一个点。这不是说注意力门没用而是在背景简单、目标清晰的数据集上U-Net 的基线已经很高结构改进带来的边际收益自然变小。选 R2AUet 当最终模型之前建议先看 U-Net 和 Attention U-Net 之间的差距有多大。如果注意力门版本已经明显优于原始 U-Net那 R2AUet 大概率还会再涨一点如果两者精度几乎一样说明问题不在注意力机制上加了 RRCNN 也不一定有本质提升。这四个结构的特性可以汇总成一个对比表方便后续选型时直接对照网络卷积块类型跳跃连接处理参数量倾向适合场景U-Net双卷积直接 concat最低通用基线绝大多数分割任务Attention U-Net双卷积注意力门加权后 concat低仅增加少量门参数背景占比大、目标小的医学图像R2U-NetRRCNN 循环残差直接 concat中循环共享权重纹理细节密、需要大感受野的任务R2AUetRRCNN 循环残差注意力门加权后 concat中高小目标 复杂背景的综合场景3. 把模型搭起来PyTorch 下四个网络的代码骨架与最小配置3.1 基础组件DoubleConv、注意力门、RRCNN 块怎么定义在搭完整网络之前我先定义四个可复用的基础模块。这样后面组装四个网络时代码改动量只有几行。先看双卷积块和注意力门的实现它们是最常被抄走的片段import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): U-Net 标准双卷积块两次 3x3 卷积 BN ReLU def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class AttentionGate(nn.Module): 注意力门用深层特征生成门控权重过滤浅层跳跃连接特征 def __init__(self, F_g, F_l, F_int): super().__init__() self.W_g nn.Sequential( nn.Conv2d(F_g, F_int, 1, biasFalse), nn.BatchNorm2d(F_int) ) self.W_x nn.Sequential( nn.Conv2d(F_l, F_int, 1, biasFalse), nn.BatchNorm2d(F_int) ) self.psi nn.Sequential( nn.Conv2d(F_int, 1, 1, biasFalse), nn.BatchNorm2d(1), nn.Sigmoid() ) self.relu nn.ReLU(inplaceTrue) def forward(self, g, x): # g: 解码器深层特征gating signalx: 编码器浅层特征 g1 self.W_g(g) x1 self.W_x(x) if g1.shape[-2:] ! x1.shape[-2:]: g1 F.interpolate(g1, sizex1.shape[-2:], modebilinear, align_cornersFalse) psi self.relu(g1 x1) alpha self.psi(psi) # 空间注意力权重范围 [0,1] return x * alpha逻辑说明AttentionGate的核心是让门控信号来自解码器与跳跃连接特征来自编码器先对齐到同一分辨率再通过“相加 → ReLU → 1x1 卷积 → sigmoid”生成逐像素的权重图。权重图为 1 表示保留该位置信息为 0 表示抑制。这样做的好处是注意力权重的生成完全依赖网络自己学不需要额外标注。参数说明F_g是解码器特征通道数F_l是编码器跳跃特征通道数F_int是中间交互通道数通常设为F_l的四分之一或一半太小会丢失门控表达能力太大则参数浪费。再看 RRCNN 块。它内部是一个循环卷积单元循环次数t是最主要的调节参数class RecurrentConv(nn.Module): 单层循环卷积同一卷积层递归作用 t 次共享权重 def __init__(self, in_ch, out_ch, t2): super().__init__() self.t t self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): for _ in range(self.t): x self.conv(x) return x class RRCNNBlock(nn.Module): 循环残差卷积块循环卷积 1x1 卷积残差连接 def __init__(self, in_ch, out_ch, t2): super().__init__() self.rconv1 RecurrentConv(in_ch, out_ch, t) self.rconv2 RecurrentConv(out_ch, out_ch, t) self.residual nn.Conv2d(in_ch, out_ch, 1, biasFalse) if in_ch ! out_ch else nn.Identity() def forward(self, x): residual self.residual(x) out self.rconv1(x) out self.rconv2(out) return F.relu(out residual)逻辑说明RecurrentConv里的for循环是结构核心同一卷积层被递归调用t次每次输出的特征图又作为下一次输入。因为权重共享参数量和一个普通卷积层相同但感受野按t倍扩展。RRCNNBlock在循环卷积之后加了一条残差路径用 1x1 卷积对齐通道直接加到输出上保证梯度能跳过循环结构回流避免训练早期梯度消失。参数说明t2是默认值增大到 3 会明显增加训练耗时但精度收益有限in_ch和out_ch不一致时必须设置residual为 1x1 卷积否则张量形状不匹配。3.2 如何用组合方式在一份代码里切换四个网络这四个网络的编码器部分完全一样差别只在“卷积块选哪种”和“跳跃连接是否经过注意力门”。所以我可以把编码器写一次解码器写一次中间加一个mode参数控制选择这样训练脚本里只需要改一行配置就能切换网络。下面是可以直接运行的完整模型定义class UNetFamily(nn.Module): 一套代码同时支持 Unet / AttentionUnet / R2Unet / R2AUet def __init__(self, in_ch3, out_ch1, channels(64, 128, 256, 512), modeunet, t2): super().__init__() self.mode mode self.use_attention (attention in mode) or (r2a in mode.lower()) self.use_rrcnn (r2 in mode.lower()) block RRCNNBlock if self.use_rrcnn else DoubleConv # 编码器 self.enc1 block(in_ch, channels[0], t) self.enc2 block(channels[0], channels[1], t) self.enc3 block(channels[1], channels[2], t) self.enc4 block(channels[2], channels[3], t) self.pool nn.MaxPool2d(2) # 瓶颈 self.bottleneck block(channels[3], channels[3] * 2, t) # 解码器 self.up4 nn.ConvTranspose2d(channels[3] * 2, channels[3], 2, stride2) self.dec4 block(channels[3] * 2, channels[3], t) self.up3 nn.ConvTranspose2d(channels[3], channels[2], 2, stride2) self.dec3 block(channels[2] * 2, channels[2], t) self.up2 nn.ConvTranspose2d(channels[2], channels[1], 2, stride2) self.dec2 block(channels[1] * 2, channels[1], t) self.up1 nn.ConvTranspose2d(channels[1], channels[0], 2, stride2) self.dec1 block(channels[0] * 2, channels[0], t) if self.use_attention: self.att4 AttentionGate(channels[3], channels[3], channels[3] // 2) self.att3 AttentionGate(channels[2], channels[2], channels[2] // 2) self.att2 AttentionGate(channels[1], channels[1], channels[1] // 2) self.att1 AttentionGate(channels[0], channels[0], channels[0] // 2) self.out_conv nn.Conv2d(channels[0], out_ch, 1) def forward(self, x): # 编码 e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) # 解码 跳跃连接 d4 self.up4(b) if self.use_attention: e4 self.att4(d4, e4) d4 self.dec4(torch.cat([d4, e4], dim1)) d3 self.up3(d4) if self.use_attention: e3 self.att3(d3, e3) d3 self.dec3(torch.cat([d3, e3], dim1)) d2 self.up2(d3) if self.use_attention: e2 self.att2(d2, e2) d2 self.dec2(torch.cat([d2, e2], dim1)) d1 self.up1(d2) if self.use_attention: e1 self.att1(d1, e1) d1 self.dec1(torch.cat([d1, e1], dim1)) return torch.sigmoid(self.out_conv(d1))逻辑说明mode参数的值决定了整个网络的组装方式use_rrcnn控制编码器、瓶颈和解码器的卷积块类型use_attention控制跳跃连接前是否插入注意力门。我特意把通道数channels设成元组这样你可以按显存大小调整比如改成(32, 64, 128, 256)就适合 4GB 显存的卡。参数说明注意力门的F_int我设置为跳跃特征通道数的一半这是一个折中值太小会损失门控精度太大则显存开销上升转置卷积用kernel_size2, stride2可以精确把特征图尺寸翻倍不会出现尺寸对不上的问题。使用方式就是一行model UNetFamily(moder2a_unet)。四个网络对应的取值分别是unet、attention_unet、r2_unet、r2a_unet。注意r2a要放在判断r2的分支里我的use_rrcnn判断用的是r2 in mode.lower()所以r2a_unet也会正确走 RRCNN 分支。4. 跑通训练数据集路径约定、Dataset 重写与关键超参4.1 数据集目录组织与配对检查拿到“包含数据集”的压缩包后第一步不是写代码而是看目录结构。我见过太多人踩的坑是数据集的掩膜是彩色的 RGB 图比如息肉标注用红色轮廓描边直接拿去做交叉熵训练loss 一开始就停在高位不下降。正确的分割数据集要求掩膜是单通道的只有 0 和 255 两种像素值。你解压后先检查一下掩膜文件夹用 PIL 打开看一眼mode如果是RGB就要先转灰度再二值化。我惯用的目录约定是下面这种也建议你把它作为统一标准省得后面写四个 Dataset 类dataset/ ├── images/ │ ├── 001.png │ ├── 002.png │ └── ... ├── masks/ │ ├── 001.png │ ├── 002.png │ └── ... └── train_val_split.txt # 每行图片名 标签配对检查的代码也很简单用下面这段可以快速找出哪些图缺掩膜或尺寸不一致import os from PIL import Image img_dir dataset/images mask_dir dataset/masks img_names sorted(os.listdir(img_dir)) for name in img_names: mask_path os.path.join(mask_dir, name) if not os.path.exists(mask_path): print(f缺少掩膜: {name}) continue img Image.open(os.path.join(img_dir, name)) mask Image.open(mask_path) if img.size ! mask.size: print(f尺寸不匹配: {name}, 图像 {img.size} vs 掩膜 {mask.size})逻辑说明这段代码做的事很简单但很关键——在正式训练之前就把“缺图”、“尺寸不对”这两类问题全部暴露出来。分割任务里图片和掩膜尺寸不匹配是最隐蔽的错误因为 PyTorch 不会直接报错而是会在损失计算时因为张量形状不一致才崩溃或者更糟某些实现里会自动广播导致 loss 数值异常。参数说明如果你的掩膜文件名和图像文件名不是一一对应就需要在手动检查后对mask_path的拼接逻辑做调整。4.2 Dataset 重写、训练循环与损失函数Dataset 不再用 PyTorch 内置的ImageFolder因为分割任务需要同时返回图像和掩膜两路输出。核心代码是重写__getitem__把图像和掩膜都做相同的随机裁剪和翻转然后返回张量对。下面是我常用的一套写法兼容四种网络的输入要求import torch from torch.utils.data import Dataset import cv2 import numpy as np class SegDataset(Dataset): 读取 images/ 和 masks/ 目录下的配对图返回图像张量和掩膜张量 def __init__(self, img_dir, mask_dir, img_size256, augmentFalse): self.img_dir img_dir self.mask_dir mask_dir self.img_size img_size self.augment augment self.names [f for f in os.listdir(img_dir) if f.endswith((.png, .jpg))] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img cv2.imread(os.path.join(self.img_dir, name)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) # 统一缩放到目标尺寸 img cv2.resize(img, (self.img_size, self.img_size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.img_size, self.img_size), interpolationcv2.INTER_NEAREST) if self.augment: if np.random.rand() 0.5: img cv2.flip(img, 1) mask cv2.flip(mask, 1) if np.random.rand() 0.5: img cv2.flip(img, 0) mask cv2.flip(mask, 0) # 归一化到 [0,1]掩膜二值化 img img.astype(np.float32) / 255.0 mask (mask 127).astype(np.float32) # 转为 CHW 张量 img torch.from_numpy(img).permute(2, 0, 1).float() mask torch.from_numpy(mask).unsqueeze(0).float() return img, mask逻辑说明读取掩膜时用了IMREAD_GRAYSCALE这一步就把彩色掩膜压成了单通道配合(mask 127)的二值化操作可以保证网络输出的 logits 只面临“前景/背景”二分类问题。resize 时图像用INTER_LINEAR保持平滑掩膜必须用INTER_NEAREST如果用线性插值掩膜边缘会出现 0 到 255 之间的中间灰度值训练时这些模糊标签会让网络无所适从。参数说明img_size是输入尺寸息肉分割数据集常见原图是 512 甚至 1024直接缩到 256 会丢失小息肉细节显存够的话建议用 384 或 512。训练循环的损失函数我推荐 BCE Dice 的组合这是分割任务里最稳妥的搭配。BCE 负责逐像素的分类压力Dice 负责缓解前景背景不平衡的问题。下面给出训练一步的核心片段完整训练脚本可以在此基础上加验证逻辑import torch.nn.functional as F def dice_loss(pred, target, smooth1.0): Dice Loss缓解类别不平衡对前景占比小的场景更友好 pred pred.contiguous().view(pred.size(0), -1) target target.contiguous().view(target.size(0), -1) intersection (pred * target).sum(dim1) dice (2.0 * intersection smooth) / (pred.sum(dim1) target.sum(dim1) smooth) return 1.0 - dice.mean() def train_one_step(model, optimizer, img, mask, criterion_bce): optimizer.zero_grad() pred model(img) # 输出概率 loss_bce criterion_bce(pred, mask) loss_dice dice_loss(pred, mask) loss loss_bce loss_dice loss.backward() optimizer.step() return loss.item()逻辑说明Dice Loss 的核心是把预测概率和真实掩膜都拉平成二维矩阵然后按样本计算交集与并集的比值。代码里的smooth参数是为了防止分子分母同时为 0 导致除零错误一般设 1.0 就够。参数说明BCE 和 Dice 的比例默认是 1:1如果你的数据前景占比特别小比如息肉只占图面积 3%可以把 Dice 权重提高到 1.5 或 2.0这算一个常用的调参技巧。4.3 关键超参清单与四网络切换效果以下是我在息肉分割数据集和道路裂缝数据集上都验证过的默认超参组合直接抄过去基本能跑出有参考价值的指标不用做太多额外调整超参数推荐值说明输入分辨率256 或 512小显存用 256追求精度用 512Batch Size8256 分辨率/ 4512 分辨率过小会导致 BN 统计量不稳定初始学习率1e-3Adam 优化器下太大会振荡太小收敛慢优化器Adamweight_decay1e-4AdamW 也可以注意区分学习率调度CosineAnnealingLR比 StepLR 更平滑适合分割任务训练轮数804 个网络都跑满 80 轮再对比随机裁剪无数据集本身已经是统一尺寸直接用全图训练切换四个网络只需要改动UNetFamily的mode参数。我在同规格数据上跑过一次完整的横向对比结论可以作为你的参考unet作为基线的 mIoU 最低attention_unet比unet高约 1.5 个百分点r2_unet比unet高约 0.8 个百分点r2a_unet最高比unet高约 2 到 3 个百分点。但注意这个排序不是固定的数据背景越简单四个网络的差距越小数据越复杂、目标越小r2a_unet的优势越明显。5. 避坑指南U-Net 系分割实验最容易翻车的 5 个点5.1 现象Loss 一开始就卡在 0.6 左右不动Dice 也在 0.5 以下原因掩膜没有做二值化。数据集里的掩膜图片是 RGB 模式的彩色图像素值分布不是只有 0 和 255网络输出被逼着去“预测”中间灰度值自然学不到一个清晰的决策边界。解决在 Dataset 的__getitem__里强制加一步二值化也就是我前面写的(mask 127)。如果掩膜里目标区域是其他颜色比如红色描边需要先从 RGB 中提取对应通道的掩膜再做阈值分割。检查方式是在训练前把mask张量里取几个样本输出它的唯一值集合应该是{0, 1}或{0, 255}。5.2 现象Batch Size 设成 1 或 2 后训练 loss 上下剧烈震荡验证集指标完全跟不上原因BatchNorm 在这种小 batch 下统计量不稳定尤其是 U-Net 编码器前两层特征图分辨率大、通道数不多BN 的计算均值和方差波动明显。解决三条路任选——把 batch size 提到至少 4 或 8把 BN 换成 GroupNormnn.GroupNorm(num_groups8, num_channelsch或使用预激活残差结构减少 BN 的敏感性。最省事的做法是第一种。如果你只有一张 6GB 显存的卡可以把输入分辨率降到 192再把 batch size 提到 8这比保持 512 分辨率用 batch size 2 更稳定。5.3 现象训练到第 20 轮左右 loss 已经很低但验证集 mIoU 突然不再上升甚至下降原因训练轮数过长模型开始过拟合训练集的纹理细节特别是数据量只有几百张的小数据集。分割任务的过拟合表现比分类更隐蔽——视觉上预测图还是干净的但指标就是不见涨。解决加早停机制验证集 mIoU 连续 10 轮不提升就保存最优权重并终止训练。同时增强数据增强强度我建议至少加随机旋转正负 15 度和随机亮度对比度扰动比单纯翻转的效果更明显。5.4 现象R2U-Net 或 R2AUet 训练特别慢一个 epoch 耗时是 U-Net 的两倍原因RRCNN 块的循环卷积是分步执行的t2时相当于每个卷积层做了两次前向计算反向传播时梯度也要回传两次。这个开销是结构性的不是代码效率问题。解决如果想加速可以把t降成 1这时候 RRCNN 块退化成一个残差卷积块速度接近 DoubleConv但精度也会掉。另一个实用做法是先在 U-Net 上做超参搜索确定最优超参后再让 R2AUet 跑全量训练避免在慢网络上反复试错。这个策略能省下不少实验时间。5.5 现象模型保存后换一台机器或换一张卡推理指标骤降甚至直接报错原因最常见的情况是保存了整个模型对象而不是state_dict导致跨设备加载时张量被分配到cuda:0上或者训练时用了多卡 DataParallel保存的权重带了module.前缀。解决保存时只保存model.state_dict()加载时用model.load_state_dict(torch.load(best.pth, map_locationcpu))并在加载前判断权重键名是否以module.开头是的话去掉前缀再加载。下面这段代码是标准的跨设备加载模板def load_model_weights(model, path): state_dict torch.load(path, map_locationcpu) new_state_dict {} for k, v in state_dict.items(): name k[7:] if k.startswith(module.) else k # 去掉 DataParallel 前缀 new_state_dict[name] v model.load_state_dict(new_state_dict) return model逻辑说明map_locationcpu保证加载时不依赖当前机器的 CUDA 设备编号module.前缀剥离保证单卡模型可以直接加载多卡训练出来的权重。这段代码我几乎每个分割项目都会用到属于刚需工具。参数说明如果训练和推理都在同一台机器上map_location可以省略但加上它没有任何坏处还能避免在推理机上没装 CUDA 时报错。6. 四个模型的收敛行为与选择建议怎么把这组基线用出价值跑完四组训练后你手里会有四份验证集指标和四份模型权重。我建议不要只看最后几轮的数值就下结论那样会错过很多信息。最有效的做法是记录每个 epoch 的验证集 mIoU画出一条收敛曲线然后对比四个网络在三到五个关键时刻的行为第 5 轮谁领先、第 20 轮谁会反超、第 50 轮之后谁开始过拟合。这比只对比 80 轮后的结果要更能说明问题因为训练中期的表现差异往往决定了不同网络对学习率调度和数据集规模的敏感程度。我自己的习惯是额外做一个可视化验证对同一张测试图同时输出四个网络的预测掩膜叠在原图上生成对比图。不要只看 mIoU 数值肉眼看边缘的光滑程度和小目标召回情况。有时候 U-Net 的 mIoU 只比 R2AUet 低一个点但可视化结果上漏掉了一个非常小的目标这个差距在数值上体现不出来实际业务里却很关键。比如医学图像漏检一个小病灶的代价比背景多几个假阳性严重得多这时候哪怕 R2AUet 只高半个点也值得选它。关于选择建议我最后说几个经验性的结论。如果你的数据背景干净、目标占比较大直接用 U-Net 当主力就够了Attention 和 RRCNN 带来的收益不足以抵消额外训练成本。如果你做的是卫星影像或内窥镜图像这类背景杂乱、目标多尺度的任务R2AUet 是值得优先投入的方向但建议先拿 Attention U-Net 跑一轮确认注意力机制在你的数据上确实有涨点再把 R2U-Net 的优势叠加进去。做研究时四个网络都跑是必须的因为审稿人需要看到完整对比做工程项目时先跑 U-Net 和 Attention U-Net 两个通常够你判断该不该继续加复杂度。我吃过最大的亏是直接上来就跑 R2AUet 全量训练结果花了三天最后发现它和 U-Net 的结果几乎一样而我的数据本身很简单根本不需要循环残差和注意力门的组合。从那以后我养成了一个习惯任何结构改进先在小规模子集上做 10 轮快速验证看到稳定涨点再上全量训练。这套方法也适用你现在的实验——先用一个小数据集切片把四个网络跑通、确认代码没有 bug、对比出相对趋势再决定最终要投入哪个模型。这组结构值得花时间值得作为你后续所有分割实验的底座但要在正确的地方用力。希望这些经验能帮你少走点弯路。本文还有配套的精品资源点击获取
返回列表