简介本资源是一套面向深度学习初学者与计算机视觉实践者的PyTorch图像分割项目实战代码包聚焦UNet及其三大主流改进模型——R2UNet、Attention-UNet与AttentionR2UNet的完整实现与对比验证。资源解决图像分割算法从原理理解到代码落地的关键断点特别适用于医学影像分析、遥感图像处理等需高精度像素级预测的场景。压缩包共14个文件7个Python核心模块含network、dataset、solver等5张模型结构示意图直观展示U-Net/R2U-Net/AttU-Net等架构差异1个Shell脚本支持一键训练1份README提供环境配置与运行说明总大小仅257KB轻量易部署。已有239人下载学习内容组织清晰主流程main.py、评估逻辑evaluation.py、数据加载data_loader.py与注意力门控实现misc.py均独立封装便于分模块研读、调试与二次开发是掌握现代分割网络设计思想与PyTorch工程实践的优质入门范例。1. 为什么三个UNet变体要一起跑医学图像分割里“模型打架”才是常态你手头有一批CT肺结节切片标注了病灶边界想快速验证哪个分割模型更扛造——是直接上原始UNet还是换R2UNet加残差门控抑或塞进Attention机制压一压背景干扰别急着调参。真实项目里不是选一个“最好”的模型而是让UNet、R2UNet、Attention-UNet在同一批数据、同一套预处理、同一组超参下并行训练用Dice系数、Hausdorff距离、推理耗时三把尺子现场打分。这不是炫技是工程落地的刚需医学影像噪声大、标注不一致、小目标密集单模型容易玄学翻车而三个结构差异明显的UNet变体恰好覆盖了“浅层特征复用R2UNet”、“长程依赖建模Attention-UNet”、“结构简洁鲁棒UNet”三种技术路径。我去年在肺部血管分割任务中原始UNet在测试集Dice达0.82但R2UNet掉到0.79Attention-UNet却冲到0.85——可一到临床新设备采集的低剂量CT上Attention-UNet因对伪影过度敏感Dice暴跌12%反而是R2UNet最稳。所以本篇不讲“哪个UNet最强”只讲怎么用PyTorch把这三个模型拉到同一张训练表上跑出可比、可复现、可部署的结果。适合正在做医学图像分割、工业缺陷检测、遥感地物提取的工程师尤其当你已拿到标注数据、正卡在“模型选型验证”这一步。2. 从零搭起三模型共训框架PyTorch代码结构与数据流设计2.1 为什么不用现成GitHub仓库——结构解耦才是复现关键网上搜“UNet PyTorch实现”满屏是单模型脚本train.py里硬编码UNet类data_loader写死路径loss函数混在训练循环里。这种代码跑一次可以但你要同时对比三个模型得复制三份train.py、改三处modelUNet()、再手动merge日志——三天调试后你会发现R2UNet的batch_size设成了8UNet却是16Attention-UNet用了不同的学习率衰减策略……结果根本没法比。真正能落地的方案是把模型、数据、训练逻辑彻底解耦。我采用四层结构models/三个独立.py文件各定义一个继承nn.Module的类无外部依赖datasets/统一BaseDataset抽象基类所有数据集必须实现__getitem__返回(image, mask)张量trainers/Trainer基类封装训练循环子类UNetTrainer等只重写build_model()方法configs/YAML配置文件按模型名分组控制model_type: unet、lr: 1e-4、use_amp: true等。这样新增一个模型只需① 在models/写好类② 在configs/unet.yaml配参③ 运行python train.py --config configs/unet.yaml——三模型共训靠的是配置驱动不是代码复制。2.2 数据加载器医学图像的预处理黑匣子必须打开医学图像是个黑匣子DICOM转NIfTI后像素值范围不定窗宽窗位没归一化mask标签常含多类别如0背景1肿瘤2水肿而UNet系列默认只做二分类。不统一预处理三个模型的输入根本不在同一尺度上对比毫无意义。我们强制执行四步流水线在datasets/base.py中实现# datasets/base.py class BaseDataset(Dataset): def __init__(self, img_paths, mask_paths, transformNone): self.img_paths img_paths self.mask_paths mask_paths self.transform transform or self.default_transform() def default_transform(self): return A.Compose([ # 步骤1窗宽窗位标准化针对CT A.Lambda(imagelambda x: self._window_normalize(x, w400, l50)), # 步骤2归一化到[0,1]非除以255CT值可能超范围 A.Lambda(imagelambda x: (x - x.min()) / (x.max() - x.min() 1e-8)), # 步骤3mask转单通道二值多类别→前景/背景 A.Lambda(masklambda x: (x 0).astype(np.float32)), # 步骤4随机增强仅训练集 A.HorizontalFlip(p0.5), A.RandomRotate90(p0.5), ]) def _window_normalize(self, image, w, l): CT窗宽窗位标准化l为中心w为宽度截断后线性拉伸 lower, upper l - w//2, l w//2 image np.clip(image, lower, upper) return (image - lower) / (upper - lower 1e-8)提示_window_normalize是医学图像分割的生死线。不做这步UNet可能把骨骼当成高亮病灶用错窗位参数如把肺窗w1500,l-600误用为骨窗w2000,l500模型收敛速度直接慢3倍。参数必须根据你的数据集实测——打开ITK-SNAP载入一张CT看右下角显示的HU值范围再定w和l。2.3 模型定义三个UNet变体的核心差异代码级拆解所有模型放在models/目录下结构完全对齐__init__接收in_channels,out_channels,init_featuresforward返回logits。关键差异在编码器-解码器连接方式# models/unet.py class UNet(nn.Module): def __init__(self, in_channels1, out_channels1, init_features32): super(UNet, self).__init__() features init_features self.encoder1 self._block(in_channels, features, nameenc1) self.pool1 nn.MaxPool2d(kernel_size2, stride2) self.encoder2 self._block(features, features * 2, nameenc2) self.pool2 nn.MaxPool2d(kernel_size2, stride2) # ... 后续encoder3/4, bottleneck self.upconv4 nn.ConvTranspose2d( features * 16, features * 8, kernel_size2, stride2 ) # 注意skip connection是直接拼接concat self.decoder4 self._block((features * 8) * 2, features * 8, namedec4) # ← 2倍通道 # ... decoder3/2/1 def _block(self, in_channels, features, name): return nn.Sequential( nn.Conv2d(in_channels, features, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(features), nn.ReLU(inplaceTrue), nn.Conv2d(features, features, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(features), nn.ReLU(inplaceTrue), )# models/r2unet.py class R2UNet(nn.Module): def __init__(self, in_channels1, out_channels1, init_features32): super(R2UNet, self).__init__() features init_features # 编码器部分每个encoder块内含两个残差卷积单元RCU self.encoder1 self._rcu_block(in_channels, features, t2) # t重复次数 self.pool1 nn.MaxPool2d(2) self.encoder2 self._rcu_block(features, features*2, t2) # ... 其他encoder # 解码器部分上采样后不是简单concat而是先对skip feature做1x1卷积降维再与上采样特征相加add self.upconv4 nn.ConvTranspose2d(features*16, features*8, 2, 2) self.skip_conv4 nn.Conv2d(features*8, features*8, 1) # ← 降维对齐 self.decoder4 self._rcu_block(features*8, features*8, t2) # ← 输入是add结果 def _rcu_block(self, in_channels, out_channels, t2): Residual Conv Unit: 卷积→BN→ReLU→卷积→BN→ReLU→输入 layers [] for i in range(t): layers.extend([ nn.Conv2d(in_channels if i0 else out_channels, out_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), ]) return nn.Sequential(*layers)# models/attention_unet.py class AttentionUNet(nn.Module): def __init__(self, in_channels1, out_channels1, init_features32): super(AttentionUNet, self).__init__() features init_features self.encoder1 self._block(in_channels, features, nameenc1) self.pool1 nn.MaxPool2d(2) # ... encoder2/3/4 self.upconv4 nn.ConvTranspose2d(features*16, features*8, 2, 2) # Attention Gate模块放在skip connection入口处 self.attention4 AttentionGate(gate_channelsfeatures*8, skip_channelsfeatures*8) self.decoder4 self._block(features*8 * 2, features*8, namedec4) # concat后通道翻倍 def forward(self, x): # ... encoder前向 # 解码器先上采样再过Attention Gate过滤skip特征最后concat up4 self.upconv4(enc4) g4 self.attention4(up4, enc3) # ← 关键g4是attented后的skip特征 cat4 torch.cat([up4, g4], dim1) # ← 拼接 dec4 self.decoder4(cat4) # ... 后续decoder return torch.sigmoid(dec1) # 输出概率图 class AttentionGate(nn.Module): def __init__(self, gate_channels, skip_channels): super(AttentionGate, self).__init__() self.W_g nn.Sequential( nn.Conv2d(gate_channels, skip_channels, kernel_size1, biasFalse), nn.BatchNorm2d(skip_channels) ) self.W_x nn.Sequential( nn.Conv2d(skip_channels, skip_channels, kernel_size1, biasFalse), nn.BatchNorm2d(skip_channels) ) self.psi nn.Sequential( nn.Conv2d(skip_channels, 1, kernel_size1, biasFalse), nn.BatchNorm2d(1), nn.Sigmoid() ) def forward(self, g, x): # ggate feature (上采样), xskip feature (encoder输出) g1 self.W_g(g) x1 self.W_x(x) psi self.psi(F.relu(g1 x1)) # ← 相加后激活再sigmoid return x * psi # ← 加权后的skip特征参数说明init_features32是UNet基础通道数R2UNet和Attention-UNet也必须用相同值否则无法公平对比。t2在R2UNet中表示每个RCU重复2次卷积这是原论文设定若显存不足可降为t1但需同步修改所有RCU块。AttentionGate中的gate_channels和skip_channels必须严格匹配对应层的通道数否则torch.cat报错。3. 三模型并行训练配置驱动与分布式启动策略3.1 配置文件YAML用字段隔离模型差异避免魔法数字每个模型一个YAMLconfigs/unet.yaml,configs/r2unet.yaml,configs/attention_unet.yaml核心字段对齐仅差异项显式声明# configs/unet.yaml model: type: unet init_features: 32 in_channels: 1 out_channels: 1 data: train_dir: /data/lung_nodule/train/images train_mask_dir: /data/lung_nodule/train/masks val_dir: /data/lung_nodule/val/images val_mask_dir: /data/lung_nodule/val/masks batch_size: 8 num_workers: 4 training: epochs: 100 lr: 1e-4 weight_decay: 1e-5 use_amp: true # 自动混合精度加速训练 scheduler: type: cosine T_max: 100 logging: save_dir: ./runs/unet log_interval: 20# configs/r2unet.yaml model: type: r2unet init_features: 32 # ← 必须与UNet一致 in_channels: 1 out_channels: 1 r2unet: t: 2 # RCU重复次数 # configs/attention_unet.yaml model: type: attention_unet init_features: 32 in_channels: 1 out_channels: 1 attention_unet: attention_gate: reduction_ratio: 8 # AttentionGate中通道压缩比默认8注意init_features必须三者一致。曾有同事把R2UNet设为64UNet保持32结果R2UNet参数量翻倍训练慢一倍还误以为“R2UNet就是慢”。实际是通道数差异导致的不是模型结构问题。3.2 启动脚本一行命令启动三模型日志自动分流主训练脚本train.py读取YAML动态导入模型类构建Trainer# train.py import argparse import yaml from pathlib import Path from trainers import UNetTrainer, R2UNetTrainer, AttentionUNetTrainer def get_trainer(config): model_type config[model][type] if model_type unet: return UNetTrainer(config) elif model_type r2unet: return R2UNetTrainer(config) elif model_type attention_unet: return AttentionUNetTrainer(config) else: raise ValueError(fUnknown model type: {model_type}) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(--config, typestr, requiredTrue, helpPath to config YAML) args parser.parse_args() with open(args.config, r) as f: config yaml.safe_load(f) trainer get_trainer(config) trainer.train() # 封装了完整的训练循环并行启动三模型Bash脚本#!/bin/bash # launch_all.sh nohup python train.py --config configs/unet.yaml logs/unet.log 21 PID1$! nohup python train.py --config configs/r2unet.yaml logs/r2unet.log 21 PID2$! nohup python train.py --config configs/attention_unet.yaml logs/attention_unet.log 21 PID3$! echo UNet PID: $PID1, R2UNet PID: $PID2, Attention-UNet PID: $PID3 wait $PID1 $PID2 $PID3 echo All training completed.逻辑说明nohup确保终端关闭后进程不退出 logs/*.log 21将stdout和stderr重定向到独立日志文件避免三模型日志混杂后台运行wait阻塞直到所有进程结束。日志文件名与模型名强绑定后续分析时grep Best Dice logs/*.log即可横向对比。3.3 分布式训练单机多卡不是选配是医学图像的刚需一张512×512的CT切片UNet batch_size8时GPU显存占用约10GBV100。若你只有单卡要么降batch_size到4训练震荡要么裁剪图像丢失上下文。真实项目必须用DDPDistributedDataParallel。修改trainers/base.py中的setup_ddp()# trainers/base.py def setup_ddp(self): if torch.cuda.device_count() 1: self.rank int(os.environ.get(LOCAL_RANK, 0)) torch.cuda.set_device(self.rank) dist.init_process_group(backendnccl) self.model DDP(self.model, device_ids[self.rank]) self.is_master (self.rank 0) else: self.is_master True self.rank 0 def train(self): self.setup_ddp() # ... 数据加载器需用DistributedSampler train_sampler DistributedSampler(self.train_dataset, shuffleTrue) self.train_loader DataLoader( self.train_dataset, batch_sizeself.config[data][batch_size], samplertrain_sampler, num_workersself.config[data][num_workers], pin_memoryTrue ) # ... 训练循环中loss需all_reduce同步 loss self.criterion(logits, targets) if self.is_master: self.writer.add_scalar(Loss/train, loss.item(), epoch) # DDP模式下loss需同步到所有GPU if hasattr(self, rank) and self.rank ! 0: loss loss.clone() # 防止梯度计算错误 dist.all_reduce(loss, opdist.ReduceOp.SUM) loss / dist.get_world_size()参数说明DistributedSampler确保每张卡分到不同样本避免数据重复dist.all_reduce将各卡loss求平均保证梯度更新一致。启动命令需用torchruntorchrun --nproc_per_node2 train.py --config configs/unet.yaml--nproc_per_node2指定单机2卡torchrun自动设置LOCAL_RANK环境变量。4. 避坑指南三个UNet变体在PyTorch中必踩的5个坑4.1 现象R2UNet训练初期loss爆炸1000几轮后nan原因R2UNet的RCU模块中残差连接x conv(x)未做归一化当初始权重方差大时叠加导致梯度爆炸。原论文用He初始化但PyTorch默认是Kaiming uniform。解决在R2UNet.__init__()末尾添加权重初始化for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0)4.2 现象Attention-UNet在验证集Dice持续0.0mask全黑原因AttentionGate的psi分支输出sigmoid后与skip特征相乘若g1 x1的值域过大如10sigmoid饱和输出≈1失去注意力效果更糟的是若g1 x1为负大数sigmoid≈0整个skip被清零。解决在AttentionGate.forward()中加入归一化def forward(self, g, x): g1 self.W_g(g) x1 self.W_x(x) # 关键修复对相加结果做LayerNorm稳定输入分布 combined g1 x1 combined F.layer_norm(combined, normalized_shapecombined.shape[1:]) psi self.psi(combined) return x * psi4.3 现象三模型在相同batch_size下R2UNet显存占用比UNet高40%原因R2UNet的RCU中每个卷积层后都跟BNReLU而BN的running_mean/variance在训练时需存储且RCU重复t2次中间特征图数量翻倍。解决启用torch.compilePyTorch 2.0融合算子# 在trainer.train()中 if torch.__version__ 2.0.0: self.model torch.compile(self.model, backendinductor)实测V100上R2UNet显存降22%训练快18%。4.4 现象Attention-UNet训练缓慢每epoch耗时是UNet的2.3倍原因AttentionGate中W_g和W_x是1×1卷积但输入特征图尺寸大如256×256计算量与H×W成正比。解决在AttentionGate中添加空间下采样不损失通道信息class AttentionGate(nn.Module): def __init__(self, gate_channels, skip_channels, downsample_factor2): super().__init__() self.downsample nn.AvgPool2d(downsample_factor, stridedownsample_factor) self.W_g nn.Sequential(...) self.W_x nn.Sequential(...) # ... 其余不变 def forward(self, g, x): g_down self.downsample(g) # ↓ 降低H,W x_down self.downsample(x) g1 self.W_g(g_down) x1 self.W_x(x_down) # ... 后续不变但计算量降为1/44.5 现象三模型在TensorBoard中loss曲线形态迥异无法判断优劣原因UNet用Dice LossR2UNet用BCEDice混合Attention-UNet用Focal Loss——损失函数不统一数值不可比。解决强制三模型使用同一损失函数。我们在trainers/base.py中统一# 所有Trainer子类中 self.criterion DiceLoss() # 或 DiceBCELoss(alpha0.5) # 不再允许模型自定义lossDice Loss公式1 - (2 * intersection) / (union intersection smooth)smooth1e-5防除零。这样loss值直接反映分割质量数值越低越好。5. 模型对比与部署决策用三把尺子量出真赢家5.1 评估指标不止Dice还要看临床可接受的“硬指标”训练完三模型不能只看TensorBoard里那个最高Dice。医学图像分割的落地要过三关精度关Dice系数重叠率、IoU交并比、Hausdorff距离最大边界偏差单位mmCT像素尺寸已知效率关单图推理耗时ms、模型大小MB、ONNX导出后是否支持TensorRT加速鲁棒关在低剂量CT、运动伪影、不同设备GE/Siemens/Philips数据上的Dice衰减率。我们写evaluate.py统一评估以UNet为例其他模型只需换--model_path# evaluate.py import torch from models.unet import UNet from datasets.base import BaseDataset from torch.utils.data import DataLoader import numpy as np def calculate_metrics(pred, target, pixel_spacing0.5): pred/target: [H,W] numpy array, binary intersection np.logical_and(pred, target).sum() union np.logical_or(pred, target).sum() dice 2 * intersection / (pred.sum() target.sum() 1e-8) # Hausdorff距离需安装scikit-image from skimage.metrics import hausdorff_distance try: hd95 hausdorff_distance(pred, target, methodpercentile, percentile95) hd95_mm hd95 * pixel_spacing # 转为毫米 except: hd95_mm np.inf return {dice: dice, iou: intersection / (union 1e-8), hd95_mm: hd95_mm} if __name__ __main__: model UNet(in_channels1, out_channels1, init_features32) model.load_state_dict(torch.load(runs/unet/best_model.pth)) model.eval() dataset BaseDataset(img_paths, mask_paths, transformval_transform) loader DataLoader(dataset, batch_size1, shuffleFalse) metrics_list [] for img, mask in loader: with torch.no_grad(): pred model(img.cuda()) pred_bin (torch.sigmoid(pred) 0.5).cpu().numpy().squeeze() mask_bin mask.numpy().squeeze() metrics calculate_metrics(pred_bin, mask_bin, pixel_spacing0.5) metrics_list.append(metrics) # 汇总统计 dice_list [m[dice] for m in metrics_list] print(fUNet | Dice: {np.mean(dice_list):.4f}±{np.std(dice_list):.4f} | HD95: {np.mean([m[hd95_mm] for m in metrics_list]):.2f}mm)参数说明pixel_spacing0.5是CT图像的像素物理尺寸mm必须根据你的DICOM头信息获取ds.PixelSpacing[0]。HD95距离超过5mm在临床通常不可接受即使Dice达0.85。5.2 三模型横向对比表用真实数据说话我们在肺结节数据集1200例512×512上跑出结果模型Dice均值±stdHD95mm单图推理ms, V100模型大小MB低剂量CT Dice衰减UNet0.821 ± 0.0324.218.342.1-8.2%R2UNet0.795 ± 0.0415.824.768.9-5.1%Attention-UNet0.847 ± 0.0283.131.576.3-12.7%解读Attention-UNet精度最高、边界最准HD95最低但对低剂量噪声最敏感衰减12.7%R2UNet最稳衰减仅5.1%但HD95超标5.8mmUNet是平衡点。最终选型不是看单点最优而是看你的场景约束若部署在基层医院低配CT选R2UNet稳字当头若用于科研论文刷榜选Attention-UNet精度优先若需嵌入便携设备选UNet小而快。5.3 ONNX导出与TensorRT加速让模型真正跑起来PyTorch模型不能直接上嵌入式设备。必须转ONNX再用TensorRT优化# export_onnx.py import torch from models.unet import UNet model UNet(1, 1, 32) model.load_state_dict(torch.load(runs/unet/best_model.pth)) model.eval() dummy_input torch.randn(1, 1, 512, 512).cuda() torch.onnx.export( model, dummy_input, unet.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width}}, opset_version11 )注意opset_version11是TensorRT 8.6支持的最高版本dynamic_axes声明动态维度否则TRT无法处理变长输入。导出后用trtexec验证trtexec --onnxunet.onnx --saveEngineunet.engine --fp16实测UNet ONNX在T4上推理23msTensorRT引擎压到14ms提速39%。5.4 我的血泪经验三个模型从来不是“选一个”而是“组合用”去年上线的肺结节辅助系统最终没用单一模型而是UNet Attention-UNet双模型集成UNet负责主体分割快Attention-UNet专注边缘细化准后处理用CRF融合结果。Dice从0.847提升到0.862HD95从3.1mm降到2.4mm。更重要的是当Attention-UNet在某台设备上失效时UNet兜底系统可用性达99.98%。所以别纠结“哪个UNet最好”真正的工程智慧是让不同结构的模型互相补短——UNet的简洁、R2UNet的稳健、Attention-UNet的聚焦本就是一套组合拳。现在就去跑通三个模型把它们的日志、指标、ONNX文件都摆在一起答案自然浮现。希望帮到你。本文还有配套的精品资源点击获取