简介这份资源面向计算机视觉方向的学习者与研究者聚焦图像分割这一核心任务提供基于Pytorch实现的UNet、R2UNet、Attention-UNet与AttentionR2UNet四种经典网络架构的完整项目代码。内容覆盖编码器-解码器结构、残差连接缓解梯度消失、注意力门控突出关键区域等关键设计适合希望从理论走向实战、系统对比不同分割模型效果的中高级读者。资源包共14个文件包含7个Python脚本、5张网络结构示意图、1个Shell运行脚本与1份说明文档压缩包约257KB代码与图示配合便于理解模型细节与训练流程。目前已有239人学习下载。读者可借助main.py、network.py、solver.py等模块快速搭建训练与评估流程结合结构图直观把握各变体的差异并在此基础上开展医学影像、自动驾驶等场景的分割实验与二次开发。1. 从一张医学影像说起UNet、R2UNet、Attention-UNet 到底该怎么选如果你手头有一批标注好的图像比如医学图像分割里的细胞切片、广告牌图像分割系统里的街景照片想用 Pytorch 跑一个能落地的分割模型大概率绕不开 UNet 这个骨架。它 2015 年提出时的 U 形结构至今没被淘汰原因很朴素浅层高分辨率特征和深层语义特征能通过跳跃连接直接拼起来小目标边缘不至于在降采样里丢干净。但真上手你会发现原始 UNet 在复杂边界、噪声多的场景下容易糊于是 R2UNet 和 Attention-UNet 这两个改进版被反复提起。这篇笔记不空谈论文我按自己复现的顺序把三个模型讲透先搭好 Pytorch 环境再把 UNet 跑通然后分别加上循环残差和注意力门控最后说清楚什么数据该用哪个、参数怎么调、哪里最容易翻车。适合已经会写基础卷积网络、想把这套东西用到自己数据集上的从业者新手跟着步骤也能走完。2. 环境搭建与数据准备别让 CUDA 和路径问题拖垮第一天2.1 Pytorch 环境搭建的稳妥路径环境这一步翻车的人最多尤其是 pytorch 安装和 cuda pytorch 下载。我的习惯是用 Anaconda 建独立环境避免和系统 Python 打架。先确认显卡驱动支持的 CUDA 版本再去 pytorch 官网拿对应命令不要凭记忆抄旧命令。# 创建独立环境python 版本选 3.9 或 3.10兼容性最好 conda create -n unet_seg python3.10 -y conda activate unet_seg # 安装 pytorch这里以 CUDA 11.8 为例具体版本以 pytorch 官网为准 # 不要直接 pip install torch容易装成 CPU 版 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 验证 GPU 是否可用 python -c import torch; print(torch.__version__, torch.cuda.is_available())逻辑说明第一条命令建环境隔离依赖第二条激活第三条是关键--index-url指向 CUDA 版 wheel 源很多人pip install torch装完发现cuda.is_available()返回 False就是源没指定对。参数上CUDA 版本必须和驱动匹配驱动版本低就选低一档的 CUDA别硬上。如果输出True环境就成了。提示安装 pytorch 是不是必须装 GPU 版本训练分割模型强烈建议 GPUCPU 跑 UNet 一个 epoch 可能要几十分钟调试成本太高。没有 N 卡就用云环境别在本地硬扛。2.2 数据集组织与标签处理unet 训练自己的数据集核心是两件事图像和掩码要一一对应掩码必须是单通道的类别索引图。我一般按下面的结构放dataset/ images/ 0001.png 0002.png masks/ 0001.png 0002.png掩码里每个像素值是类别编号二分类就是 0 和 1多分类是 0 到 N-1。常见错误是掩码存成了三通道 RGB或者用了调色板 PNG读进来值不对。用下面这段检查import numpy as np from PIL import Image mask np.array(Image.open(dataset/masks/0001.png)) print(mask.shape, np.unique(mask)) # 期望输出类似 (512, 512) [0 1]如果是 (512, 512, 3) 就要转灰度逻辑说明np.unique能直接暴露标签值是否连续、是否有多余通道。如果输出一堆 0 到 255 的值说明掩码是灰度图不是索引图需要重新映射。参数上图像尺寸不要求统一但训练时会被 resize 到固定大小建议原图短边不低于 256否则小目标 resize 后更糊。2.3 数据增强与 DataLoader 的最小实现分割任务的数据增强必须图像和掩码同步变换用 albumentations 最省事。import albumentations as A from torch.utils.data import Dataset, DataLoader import cv2 train_tf A.Compose([ A.Resize(256, 256), A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.3), ]) class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, tf): self.img_dir, self.mask_dir, self.tf img_dir, mask_dir, tf self.names os.listdir(img_dir) def __getitem__(self, i): name self.names[i] img cv2.imread(f{self.img_dir}/{name})[:, :, ::-1] # BGR 转 RGB mask cv2.imread(f{self.mask_dir}/{name}, 0) # 单通道读 aug self.tf(imageimg, maskmask) img aug[image].transpose(2, 0, 1) / 255.0 # HWC 转 CHW 并归一化 return torch.tensor(img, dtypetorch.float32), torch.tensor(aug[mask], dtypetorch.long)逻辑说明A.Compose里所有几何变换会同时作用到 image 和 mask这是分割增强的关键自己写容易漏。cv2.imread第二参数 0 保证掩码单通道。归一化除以 255 是常规操作如果用了预训练权重要换成 ImageNet 的均值和方差。dtypetorch.long是因为交叉熵损失要求标签是 long 类型这点不注意会报错。3. UNet 基线复现把跳跃连接和损失函数讲清楚3.1 UNet 结构拆解与编码器实现UNet 的 U 形由编码器、瓶颈、解码器组成编码器每层两次卷积加一次池化通道数翻倍。我一般把卷积块抽出来复用import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_c, out_c): super().__init__() self.block nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), nn.Conv2d(out_c, out_c, 3, padding1), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)逻辑说明padding1保证卷积后尺寸不变这是 UNet 能精确拼接的前提。BatchNorm 加速收敛但 batch size 很小时比如 2会不稳定可以换 GroupNorm。参数上in_c和out_c决定通道流转编码器第一层通常 64 通道起步显存不够就减半。3.2 解码器、跳跃连接与完整前向解码器上采样后和编码器对应层拼接再走一次 DoubleConv。class UNet(nn.Module): def __init__(self, in_c3, num_classes2): super().__init__() self.d1 DoubleConv(in_c, 64) self.d2 DoubleConv(64, 128) self.d3 DoubleConv(128, 256) self.d4 DoubleConv(256, 512) self.bottleneck DoubleConv(512, 1024) self.pool nn.MaxPool2d(2) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.u4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.u3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.u2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.u1 DoubleConv(128, 64) self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): c1 self.d1(x) c2 self.d2(self.pool(c1)) c3 self.d3(self.pool(c2)) c4 self.d4(self.pool(c3)) b self.bottleneck(self.pool(c4)) x self.u4(torch.cat([self.up4(b), c4], dim1)) x self.u3(torch.cat([self.up3(x), c3], dim1)) x self.u2(torch.cat([self.up2(x), c2], dim1)) x self.u1(torch.cat([self.up1(x), c1], dim1)) return self.out(x)逻辑说明torch.cat沿通道维拼接所以up4输出 512 通道和c4的 512 拼成 1024正好喂给u4。这是 UNet 最核心的设计浅层细节和深层语义在这里汇合。ConvTranspose2d的stride2实现两倍上采样也可以用nn.Upsample加卷积替代前者有可学习参数后者更省显存。输出层 1x1 卷积把通道压到类别数。3.3 损失函数与训练循环分割最常用交叉熵加 Dice 的组合前者管像素分类后者管区域重叠。def dice_loss(pred, target, eps1e-6): pred torch.softmax(pred, dim1)[:, 1] # 取前景概率 target (target 1).float() inter (pred * target).sum() return 1 - (2 * inter eps) / (pred.sum() target.sum() eps) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) for epoch in range(50): model.train() for img, mask in loader: img, mask img.cuda(), mask.cuda() pred model(img) loss criterion(pred, mask) dice_loss(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step()逻辑说明dice_loss里先 softmax 取前景通道再算交并比eps防止除零。总损失是两者相加权重可以调边界差的数据集把 Dice 权重加大。学习率 1e-3 是 Adam 的常规起点loss 震荡就降到 1e-4。zero_grad必须在 backward 前顺序错了梯度会累积。注意unet 使用时的注意事项里最常见的是输入尺寸不能被 16 整除。UNet 下采样 4 次尺寸必须是 16 的倍数否则拼接时尺寸对不上。训练前统一 resize 到 256 或 512 最省心。4. R2UNet 与 Attention-UNet两个改进到底改在哪4.1 R2UNet 的循环残差块实现R2UNet 的核心是 Recurrent Residual Block把普通卷积换成带循环的残差单元让特征在块内反复 refinement。它适合边界模糊、需要多轮修正的医学图像分割。class RRBlock(nn.Module): def __init__(self, in_c, out_c, t2): super().__init__() self.t t self.conv nn.Conv2d(in_c, out_c, 3, padding1) self.bn nn.BatchNorm2d(out_c) self.relu nn.ReLU(inplaceTrue) def forward(self, x): # 循环 t 次每次把输入和输出相加再激活 for _ in range(self.t): x self.relu(self.bn(self.conv(x)) x) return x逻辑说明t是循环次数论文里常用 2。每次循环把卷积结果和原输入相加形成残差再激活。这样同一个块内特征被反复处理感受野不变但表达更强。参数上t越大越强也越慢2 是精度和速度的平衡点。把 UNet 里的 DoubleConv 换成 RRBlock其余结构不动就是 R2UNet。4.2 Attention Gate 的加法和位置Attention-UNet 在跳跃连接上加注意力门控让解码器只关注编码器里和当前区域相关的特征抑制背景干扰。适合前景小、背景杂的场景比如广告牌图像分割系统里从街景中抠出牌子。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), nn.BatchNorm2d(F_int)) self.W_x nn.Sequential(nn.Conv2d(F_l, F_int, 1), nn.BatchNorm2d(F_int)) self.psi nn.Sequential(nn.Conv2d(F_int, 1, 1), nn.BatchNorm2d(1), nn.Sigmoid()) def forward(self, g, x): # g 是解码器上采样特征x 是编码器跳跃特征 att self.psi(torch.relu(self.W_g(g) self.W_x(x))) return x * att逻辑说明W_g和W_x把两路特征压到同一中间通道F_int相加后过 sigmoid 得到 0 到 1 的注意力图再乘回编码器特征。F_int一般取F_l的一半。加的位置在torch.cat之前先对c4做门控再拼接。参数上注意力图是逐像素的所以对小目标特别有效但会引入额外计算显存紧张时减少F_int。4.3 三个模型的训练对比与选型模型参数量趋势适合场景训练速度UNet基准通用、边界清晰最快R2UNet略增医学图像、边界模糊较慢Attention-UNet中等增加小目标、背景复杂中等选型上我的经验数据量小、边界清楚UNet 就够别过度设计医学图像分割这种边界渐变、需要反复修正的上 R2UNet前景占比小、背景干扰大的Attention-UNet 收益最明显。三个模型可以共用同一套训练循环只换模型类方便对比。5. 避坑与排查那些让我重跑一整晚的问题5.1 损失不下降先查标签和归一化现象训练几个 epochloss 卡在 0.69 附近不动正好是二分类交叉熵的初始值。原因通常是标签全是 0或者掩码读进来值域不对。解决用np.unique打印标签分布确认前景存在检查归一化是否把图像压到了合理范围除以 255 后应该在 0 到 1。5.2 显存溢出从 batch size 和尺寸下手现象CUDA out of memory尤其跑 R2UNet 时。原因是循环块和注意力门控都增加显存。解决先把 batch size 降到 2 或 1再把输入尺寸从 512 降到 256最后考虑混合精度训练。别一上来就换大显卡参数没调对换卡也白搭。5.3 预测结果全黑或全白现象推理输出要么全背景要么全前景。原因是模型没收敛或者推理时忘了model.eval()和torch.no_grad()BatchNorm 用了训练时的统计量。解决推理前加model.eval()用with torch.no_grad():包住前向输出取 argmax 得到类别图。5.4 拼接维度报错现象torch.cat报通道数不匹配。原因是输入尺寸不是 16 的倍数下采样后某层尺寸和上采样对不上。解决训练和推理统一 resize 到 16 的倍数比如 256、512。如果必须用原始尺寸用 padding 补齐到最近的 16 倍数推理后再裁回去。5.5 验证集指标虚高现象Dice 很高但可视化一看全是糊的。原因是前景占比极低时全预测背景也能有高准确率Dice 对类别不平衡敏感。解决同时看 IoU 和可视化结果别只信一个数字。数据增强里加随机裁剪让前景占比更均衡。6. 进阶技巧把训练好的模型导出 ONNX 并验证一致性模型训完只是第一步落地往往要转 ONNX 部署。pytorch 转 onnx 有个坑动态尺寸和 BatchNorm 的处理。我一般固定输入尺寸导出再逐层对比输出。import torch model.eval() dummy torch.randn(1, 3, 256, 256).cuda() torch.onnx.export( model, dummy, unet.onnx, input_names[input], output_names[output], opset_version11, # 11 兼容性好别盲目追高 dynamic_axesNone, # 固定尺寸避免动态轴带来的算子问题 ) # 验证 ONNX 和 Pytorch 输出一致性 import onnxruntime as ort import numpy as np sess ort.InferenceSession(unet.onnx, providers[CUDAExecutionProvider]) with torch.no_grad(): pt_out model(dummy).cpu().numpy() onnx_out sess.run(None, {input: dummy.cpu().numpy()})[0] print(np.abs(pt_out - onnx_out).max()) # 一般应小于 1e-4逻辑说明opset_version11是分割模型常用的稳定版本太高有些算子不支持。dynamic_axesNone固定尺寸部署时输入必须一致否则要重新导出。一致性检查里最大绝对误差小于 1e-4 说明转换没问题超过就要查是哪一层算子对不上通常是上采样或 BatchNorm。参数上如果部署环境只支持 CPU把 provider 换成CPUExecutionProvider速度会慢但结果一致。我自己的习惯是每次改完模型结构先跑一遍 ONNX 一致性验证再上训练能提前发现算子兼容问题省得训完才发现导不出去。这套 UNet 加两个改进版的组合我前后在几个项目里用过最深的教训是别一上来就堆改进先把基线跑稳、把数据和损失理清楚改进才有意义。希望帮到你。本文还有配套的精品资源点击获取 Pytorch图像分割实战：UNet、R2UNet与Attention-UNet选型与避坑指南
简介这份资源面向计算机视觉方向的学习者与研究者聚焦图像分割这一核心任务提供基于Pytorch实现的UNet、R2UNet、Attention-UNet与AttentionR2UNet四种经典网络架构的完整项目代码。内容覆盖编码器-解码器结构、残差连接缓解梯度消失、注意力门控突出关键区域等关键设计适合希望从理论走向实战、系统对比不同分割模型效果的中高级读者。资源包共14个文件包含7个Python脚本、5张网络结构示意图、1个Shell运行脚本与1份说明文档压缩包约257KB代码与图示配合便于理解模型细节与训练流程。目前已有239人学习下载。读者可借助main.py、network.py、solver.py等模块快速搭建训练与评估流程结合结构图直观把握各变体的差异并在此基础上开展医学影像、自动驾驶等场景的分割实验与二次开发。1. 从一张医学影像说起UNet、R2UNet、Attention-UNet 到底该怎么选如果你手头有一批标注好的图像比如医学图像分割里的细胞切片、广告牌图像分割系统里的街景照片想用 Pytorch 跑一个能落地的分割模型大概率绕不开 UNet 这个骨架。它 2015 年提出时的 U 形结构至今没被淘汰原因很朴素浅层高分辨率特征和深层语义特征能通过跳跃连接直接拼起来小目标边缘不至于在降采样里丢干净。但真上手你会发现原始 UNet 在复杂边界、噪声多的场景下容易糊于是 R2UNet 和 Attention-UNet 这两个改进版被反复提起。这篇笔记不空谈论文我按自己复现的顺序把三个模型讲透先搭好 Pytorch 环境再把 UNet 跑通然后分别加上循环残差和注意力门控最后说清楚什么数据该用哪个、参数怎么调、哪里最容易翻车。适合已经会写基础卷积网络、想把这套东西用到自己数据集上的从业者新手跟着步骤也能走完。2. 环境搭建与数据准备别让 CUDA 和路径问题拖垮第一天2.1 Pytorch 环境搭建的稳妥路径环境这一步翻车的人最多尤其是 pytorch 安装和 cuda pytorch 下载。我的习惯是用 Anaconda 建独立环境避免和系统 Python 打架。先确认显卡驱动支持的 CUDA 版本再去 pytorch 官网拿对应命令不要凭记忆抄旧命令。# 创建独立环境python 版本选 3.9 或 3.10兼容性最好 conda create -n unet_seg python3.10 -y conda activate unet_seg # 安装 pytorch这里以 CUDA 11.8 为例具体版本以 pytorch 官网为准 # 不要直接 pip install torch容易装成 CPU 版 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 验证 GPU 是否可用 python -c import torch; print(torch.__version__, torch.cuda.is_available())逻辑说明第一条命令建环境隔离依赖第二条激活第三条是关键--index-url指向 CUDA 版 wheel 源很多人pip install torch装完发现cuda.is_available()返回 False就是源没指定对。参数上CUDA 版本必须和驱动匹配驱动版本低就选低一档的 CUDA别硬上。如果输出True环境就成了。提示安装 pytorch 是不是必须装 GPU 版本训练分割模型强烈建议 GPUCPU 跑 UNet 一个 epoch 可能要几十分钟调试成本太高。没有 N 卡就用云环境别在本地硬扛。2.2 数据集组织与标签处理unet 训练自己的数据集核心是两件事图像和掩码要一一对应掩码必须是单通道的类别索引图。我一般按下面的结构放dataset/ images/ 0001.png 0002.png masks/ 0001.png 0002.png掩码里每个像素值是类别编号二分类就是 0 和 1多分类是 0 到 N-1。常见错误是掩码存成了三通道 RGB或者用了调色板 PNG读进来值不对。用下面这段检查import numpy as np from PIL import Image mask np.array(Image.open(dataset/masks/0001.png)) print(mask.shape, np.unique(mask)) # 期望输出类似 (512, 512) [0 1]如果是 (512, 512, 3) 就要转灰度逻辑说明np.unique能直接暴露标签值是否连续、是否有多余通道。如果输出一堆 0 到 255 的值说明掩码是灰度图不是索引图需要重新映射。参数上图像尺寸不要求统一但训练时会被 resize 到固定大小建议原图短边不低于 256否则小目标 resize 后更糊。2.3 数据增强与 DataLoader 的最小实现分割任务的数据增强必须图像和掩码同步变换用 albumentations 最省事。import albumentations as A from torch.utils.data import Dataset, DataLoader import cv2 train_tf A.Compose([ A.Resize(256, 256), A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.3), ]) class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, tf): self.img_dir, self.mask_dir, self.tf img_dir, mask_dir, tf self.names os.listdir(img_dir) def __getitem__(self, i): name self.names[i] img cv2.imread(f{self.img_dir}/{name})[:, :, ::-1] # BGR 转 RGB mask cv2.imread(f{self.mask_dir}/{name}, 0) # 单通道读 aug self.tf(imageimg, maskmask) img aug[image].transpose(2, 0, 1) / 255.0 # HWC 转 CHW 并归一化 return torch.tensor(img, dtypetorch.float32), torch.tensor(aug[mask], dtypetorch.long)逻辑说明A.Compose里所有几何变换会同时作用到 image 和 mask这是分割增强的关键自己写容易漏。cv2.imread第二参数 0 保证掩码单通道。归一化除以 255 是常规操作如果用了预训练权重要换成 ImageNet 的均值和方差。dtypetorch.long是因为交叉熵损失要求标签是 long 类型这点不注意会报错。3. UNet 基线复现把跳跃连接和损失函数讲清楚3.1 UNet 结构拆解与编码器实现UNet 的 U 形由编码器、瓶颈、解码器组成编码器每层两次卷积加一次池化通道数翻倍。我一般把卷积块抽出来复用import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_c, out_c): super().__init__() self.block nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), nn.Conv2d(out_c, out_c, 3, padding1), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)逻辑说明padding1保证卷积后尺寸不变这是 UNet 能精确拼接的前提。BatchNorm 加速收敛但 batch size 很小时比如 2会不稳定可以换 GroupNorm。参数上in_c和out_c决定通道流转编码器第一层通常 64 通道起步显存不够就减半。3.2 解码器、跳跃连接与完整前向解码器上采样后和编码器对应层拼接再走一次 DoubleConv。class UNet(nn.Module): def __init__(self, in_c3, num_classes2): super().__init__() self.d1 DoubleConv(in_c, 64) self.d2 DoubleConv(64, 128) self.d3 DoubleConv(128, 256) self.d4 DoubleConv(256, 512) self.bottleneck DoubleConv(512, 1024) self.pool nn.MaxPool2d(2) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.u4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.u3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.u2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.u1 DoubleConv(128, 64) self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): c1 self.d1(x) c2 self.d2(self.pool(c1)) c3 self.d3(self.pool(c2)) c4 self.d4(self.pool(c3)) b self.bottleneck(self.pool(c4)) x self.u4(torch.cat([self.up4(b), c4], dim1)) x self.u3(torch.cat([self.up3(x), c3], dim1)) x self.u2(torch.cat([self.up2(x), c2], dim1)) x self.u1(torch.cat([self.up1(x), c1], dim1)) return self.out(x)逻辑说明torch.cat沿通道维拼接所以up4输出 512 通道和c4的 512 拼成 1024正好喂给u4。这是 UNet 最核心的设计浅层细节和深层语义在这里汇合。ConvTranspose2d的stride2实现两倍上采样也可以用nn.Upsample加卷积替代前者有可学习参数后者更省显存。输出层 1x1 卷积把通道压到类别数。3.3 损失函数与训练循环分割最常用交叉熵加 Dice 的组合前者管像素分类后者管区域重叠。def dice_loss(pred, target, eps1e-6): pred torch.softmax(pred, dim1)[:, 1] # 取前景概率 target (target 1).float() inter (pred * target).sum() return 1 - (2 * inter eps) / (pred.sum() target.sum() eps) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) for epoch in range(50): model.train() for img, mask in loader: img, mask img.cuda(), mask.cuda() pred model(img) loss criterion(pred, mask) dice_loss(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step()逻辑说明dice_loss里先 softmax 取前景通道再算交并比eps防止除零。总损失是两者相加权重可以调边界差的数据集把 Dice 权重加大。学习率 1e-3 是 Adam 的常规起点loss 震荡就降到 1e-4。zero_grad必须在 backward 前顺序错了梯度会累积。注意unet 使用时的注意事项里最常见的是输入尺寸不能被 16 整除。UNet 下采样 4 次尺寸必须是 16 的倍数否则拼接时尺寸对不上。训练前统一 resize 到 256 或 512 最省心。4. R2UNet 与 Attention-UNet两个改进到底改在哪4.1 R2UNet 的循环残差块实现R2UNet 的核心是 Recurrent Residual Block把普通卷积换成带循环的残差单元让特征在块内反复 refinement。它适合边界模糊、需要多轮修正的医学图像分割。class RRBlock(nn.Module): def __init__(self, in_c, out_c, t2): super().__init__() self.t t self.conv nn.Conv2d(in_c, out_c, 3, padding1) self.bn nn.BatchNorm2d(out_c) self.relu nn.ReLU(inplaceTrue) def forward(self, x): # 循环 t 次每次把输入和输出相加再激活 for _ in range(self.t): x self.relu(self.bn(self.conv(x)) x) return x逻辑说明t是循环次数论文里常用 2。每次循环把卷积结果和原输入相加形成残差再激活。这样同一个块内特征被反复处理感受野不变但表达更强。参数上t越大越强也越慢2 是精度和速度的平衡点。把 UNet 里的 DoubleConv 换成 RRBlock其余结构不动就是 R2UNet。4.2 Attention Gate 的加法和位置Attention-UNet 在跳跃连接上加注意力门控让解码器只关注编码器里和当前区域相关的特征抑制背景干扰。适合前景小、背景杂的场景比如广告牌图像分割系统里从街景中抠出牌子。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), nn.BatchNorm2d(F_int)) self.W_x nn.Sequential(nn.Conv2d(F_l, F_int, 1), nn.BatchNorm2d(F_int)) self.psi nn.Sequential(nn.Conv2d(F_int, 1, 1), nn.BatchNorm2d(1), nn.Sigmoid()) def forward(self, g, x): # g 是解码器上采样特征x 是编码器跳跃特征 att self.psi(torch.relu(self.W_g(g) self.W_x(x))) return x * att逻辑说明W_g和W_x把两路特征压到同一中间通道F_int相加后过 sigmoid 得到 0 到 1 的注意力图再乘回编码器特征。F_int一般取F_l的一半。加的位置在torch.cat之前先对c4做门控再拼接。参数上注意力图是逐像素的所以对小目标特别有效但会引入额外计算显存紧张时减少F_int。4.3 三个模型的训练对比与选型模型参数量趋势适合场景训练速度UNet基准通用、边界清晰最快R2UNet略增医学图像、边界模糊较慢Attention-UNet中等增加小目标、背景复杂中等选型上我的经验数据量小、边界清楚UNet 就够别过度设计医学图像分割这种边界渐变、需要反复修正的上 R2UNet前景占比小、背景干扰大的Attention-UNet 收益最明显。三个模型可以共用同一套训练循环只换模型类方便对比。5. 避坑与排查那些让我重跑一整晚的问题5.1 损失不下降先查标签和归一化现象训练几个 epochloss 卡在 0.69 附近不动正好是二分类交叉熵的初始值。原因通常是标签全是 0或者掩码读进来值域不对。解决用np.unique打印标签分布确认前景存在检查归一化是否把图像压到了合理范围除以 255 后应该在 0 到 1。5.2 显存溢出从 batch size 和尺寸下手现象CUDA out of memory尤其跑 R2UNet 时。原因是循环块和注意力门控都增加显存。解决先把 batch size 降到 2 或 1再把输入尺寸从 512 降到 256最后考虑混合精度训练。别一上来就换大显卡参数没调对换卡也白搭。5.3 预测结果全黑或全白现象推理输出要么全背景要么全前景。原因是模型没收敛或者推理时忘了model.eval()和torch.no_grad()BatchNorm 用了训练时的统计量。解决推理前加model.eval()用with torch.no_grad():包住前向输出取 argmax 得到类别图。5.4 拼接维度报错现象torch.cat报通道数不匹配。原因是输入尺寸不是 16 的倍数下采样后某层尺寸和上采样对不上。解决训练和推理统一 resize 到 16 的倍数比如 256、512。如果必须用原始尺寸用 padding 补齐到最近的 16 倍数推理后再裁回去。5.5 验证集指标虚高现象Dice 很高但可视化一看全是糊的。原因是前景占比极低时全预测背景也能有高准确率Dice 对类别不平衡敏感。解决同时看 IoU 和可视化结果别只信一个数字。数据增强里加随机裁剪让前景占比更均衡。6. 进阶技巧把训练好的模型导出 ONNX 并验证一致性模型训完只是第一步落地往往要转 ONNX 部署。pytorch 转 onnx 有个坑动态尺寸和 BatchNorm 的处理。我一般固定输入尺寸导出再逐层对比输出。import torch model.eval() dummy torch.randn(1, 3, 256, 256).cuda() torch.onnx.export( model, dummy, unet.onnx, input_names[input], output_names[output], opset_version11, # 11 兼容性好别盲目追高 dynamic_axesNone, # 固定尺寸避免动态轴带来的算子问题 ) # 验证 ONNX 和 Pytorch 输出一致性 import onnxruntime as ort import numpy as np sess ort.InferenceSession(unet.onnx, providers[CUDAExecutionProvider]) with torch.no_grad(): pt_out model(dummy).cpu().numpy() onnx_out sess.run(None, {input: dummy.cpu().numpy()})[0] print(np.abs(pt_out - onnx_out).max()) # 一般应小于 1e-4逻辑说明opset_version11是分割模型常用的稳定版本太高有些算子不支持。dynamic_axesNone固定尺寸部署时输入必须一致否则要重新导出。一致性检查里最大绝对误差小于 1e-4 说明转换没问题超过就要查是哪一层算子对不上通常是上采样或 BatchNorm。参数上如果部署环境只支持 CPU把 provider 换成CPUExecutionProvider速度会慢但结果一致。我自己的习惯是每次改完模型结构先跑一遍 ONNX 一致性验证再上训练能提前发现算子兼容问题省得训完才发现导不出去。这套 UNet 加两个改进版的组合我前后在几个项目里用过最深的教训是别一上来就堆改进先把基线跑稳、把数据和损失理清楚改进才有意义。希望帮到你。本文还有配套的精品资源点击获取