简介这套代码面向医学图像分割研究者与开发者围绕CT等医学影像的语义分割任务完整实现U-Net与Attention U-Net双模型框架。系统覆盖数据处理、模型构建、训练评估与预测推理全流程支持自定义数据路径与格式、随机翻转增强、窗宽窗位对比度调整并通过注意力门控和跳跃连接加强特征融合适应多类别分割场景。压缩包共14个文件以5个Python源文件为主干分别对应dataset、model、train、predict、utils五大功能模块另含7个pyc预编译文件以及readme和依赖说明整体仅16KB结构清晰、便于快速阅读与二次开发。目前已有117人学习浏览。学习这套代码可直接获得可运行的分割项目训练过程自动保存最佳模型与日志并输出叠加掩码的可视化预测结果适合想在医学图像分割方向快速搭建基线或对比两种网络效果的初中级开发者。1. 为什么医学图像分割绕不开U-Net与Attention U-Net医学图像分割里有个反直觉现象不少榜单靠前的模型换一个数据集就失效而 U-Net 至今仍是标准 baseline。关键在于医学图像样本量小、标注稀疏、边界模糊带跳跃连接的编码-解码结构风险最低。Attention U-Net 的核心改动只是把注意力门控插到跳跃连接上对目标加权、对背景抑制。这套系统解决的是落地问题给定一批 CT/MRI 和手工标注 mask从数据读入、网络定义、训练到指标计算一条链路跑通。真正的坑在注意力门控的维度对齐、类别不平衡的损失、以及把注意力系数抽出来核验分割质量。适合刚入门的工程师也适合想量化 Attention U-Net 收益边界的 U-Net 老用户。接下来的内容按结构原理 → 代码实现 → 数据与损失 → 验证技巧展开中间三章是可复现的代码和参数最后一章给出排查注意力失效的具体手段。2. U-Net的编码-解码结构与Attention U-Net的注意力接入点2.1 编码路径分辨率递减与感受野扩张的分工U-Net 的编码路径由双卷积 下采样单元逐级堆叠而成。以 256×256 的单通道输入为例经过 4 次下采样特征图尺寸从 256 降到 128、64、32、16通道数从 64 逐级翻倍到 512。每级内部是两个 3×3 卷积各接 BatchNorm 和 ReLU下采样用 MaxPool 或 stride2 的卷积后者少一次池化显存上略占优两者在最终指标上差别不大。这个结构对应一个朴素观察浅层特征保存边界和纹理深层特征保存语义类别。解码部分的任务是把深层高语义、低分辨率的特征逐步恢复回原分辨率。分辨率每下降一次细节信息就丢一部分于是跳跃连接成为 U-Net 的骨架级设计而不是可选项。2.2 跳跃连接的双重作用信息补偿与梯度捷径跳跃连接把编码器第 i 层的特征直接接到解码器对应层和上采样后的特征在通道维拼接。第一个作用是信息补偿深层特征经过多次下采样后边缘被严重平滑把浅层特征拼回来等于重新注入边界信息这是医学图像分割对边缘敏感的根本保障。第二个作用容易被忽略梯度捷径。分割损失主要落在解码器的末端输出上反向传播时浅层卷积的梯度要穿过整个解码路径拼接连接给梯度开了一条短路径训练初期收敛明显更稳。这个性质在 batch size 小、梯度噪声大的时候尤其保值也是 U-Net 在小数据集上比纯 FCN 类型网络好训练的原因之一。2.3 Attention Gate加在拼接之前用加法注意力算权重Attention U-Net 的改动只有一个模块Attention Gate位置在跳跃连接特征与解码特征拼接之前。它有两个输入x 是编码器同层特征细节丰富但噪声多g 是解码器上一级上采样后的特征语义更强。AG 用 g 计算逐像素权重 alpha把 x 中与目标无关的区域按下再送入拼接。权重计算是加法注意力x 和 g 分别经 1×1 卷积映射到同一中间维度相加后过 ReLU再经 1×1 卷积压到单通道sigmoid 输出 0 到 1 的 alpha。核心逻辑是让语义特征 g 告诉细节特征 x 哪里值得保留。注意这里用的是加法融合而不是点乘点乘注意力对维度对齐敏感在特征图分辨率不一致的尺度上更容易数值不稳。代码实现import torch import torch.nn as nn import torch.nn.functional as F class AttentionGate(nn.Module): def __init__(self, in_ch_x, in_ch_g, inter_ch64): super().__init__() self.wx nn.Conv2d(in_ch_x, inter_ch, kernel_size1, biasFalse) self.wg nn.Conv2d(in_ch_g, inter_ch, kernel_size1, biasFalse) self.psi nn.Conv2d(inter_ch, 1, kernel_size1, biasTrue) self.relu nn.ReLU(inplaceTrue) self.sigmoid nn.Sigmoid() def forward(self, x, g): # x: 编码路径跳跃连接 (B, in_ch_x, H, W) # g: 解码路径上采样后的门控特征 (B, in_ch_g, H, W) if g.shape[-2:] ! x.shape[-2:]: g F.interpolate(g, sizex.shape[-2:], modebilinear, align_cornersFalse) fx self.wx(x) fg self.wg(g) a self.psi(self.relu(fx fg)) alpha self.sigmoid(a) return alpha * x, alpha逻辑说明wx 和 wg 分别把 x、g 投影到 inter_ch 维两者相加让两个来源的特征在同一个空间中互相调制psi 再把调制结果压缩到单通道。返回的第一个值是加权后的跳跃特征第二个是注意力系数后者在验证阶段会被抽出来做可视化。参数说明in_ch_x 是当前尺度编码特征通道数in_ch_g 是解码特征通道数inter_ch 论文默认 64全部用 1×1 卷积因此这个模块可以在任意分辨率下复用。2.4 加入AG对系统的代价与收益边界把 U-Net 和 Attention U-Net 放进同一套分割系统差异集中在跳跃连接的处理方式和新增参数上对比项U-NetAttention U-Net跳跃连接直接拼接权重靠后续卷积学习先乘注意力系数再拼接新增参数无每层 AG 约 3 个 1×1 卷积总量几十万推理速度基准根据 AG 数量慢 5%~10%低对比度小病灶容易漏检漏检率下降但目标邻域也可能被误抑制训练初期正常alpha 可能饱和在 1 附近需要观察注意AG 不是每层都加收益最大。原论文加在较浅的尺度上最深层接的是纯语义信息往往不需要门控。实际项目里如果病灶只占图像 1% 以下加在 256、128 分辨率的 AG 反而容易把细小目标整体压掉我一般只在 1/4 和 1/8 分辨率的两层加。训练初期的 alpha 饱和是 Attention U-Net 最常见的失败模式sigmoid 在输入为 0 附近输出正好是 0.5但经过 ReLU 和残差式的前向叠加第一代参数很可能把 alpha 推到接近 1让门控退化成恒等映射。判断方法很简单前 10 个 epoch 打印 alpha 的均值如果一直高于 0.95把学习率降到 5e-5 重训通常比调网络结构更有效。3. 用PyTorch搭建Attention U-Net的最小可复现分割代码3.1 双卷积与下采样把编码器写成可复用单元整个模型可以收敛成四个部件DoubleConv、Down、Up、AttentionGate。先写前两个编码路径的每个下采样阶段都是同一套逻辑只是通道数不同class DoubleConv(nn.Module): 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 Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool nn.MaxPool2d(2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x): return self.conv(self.pool(x))逻辑说明DoubleConv 是两个 3×3 卷积加 BN 加 ReLUpadding1 保证特征图尺寸不变这样尺寸变化只由 Down 中的 MaxPool2d(2) 控制。参数说明biasFalse 是因为卷积后面接 BatchNormBN 层自带可学习的偏置卷积层再带 bias 是冗余的还会轻微影响 BN 的统计估计。inplaceTrue省一块显存实测在 256 输入下每层能省几 MB对大 batch 有帮助。3.2 解码路径接线顺序先上采样再过AG再拼接解码器的顺序有讲究先对深层特征做转置卷积上采样让解码特征与跳跃特征空间尺寸一致再送入 AttentionGate最后拼接。顺序反了会导致 AG 的分辨率检查分支总是被触发虽然功能正确但 interpolate 会引入额外的数值误差class Up(nn.Module): def __init__(self, in_ch, skip_ch, out_ch, use_attTrue): super().__init__() self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2) self.ag AttentionGate(skip_ch, in_ch // 2) if use_att else None self.conv DoubleConv(in_ch // 2 skip_ch, out_ch) def forward(self, x, skip): x self.up(x) # 先恢复分辨率 if self.ag is not None: skip, _ self.ag(skip, x) # 门控作用于跳跃连接 x torch.cat([x, skip], dim1) return self.conv(x)逻辑说明ConvTranspose2d 把解码特征上采样一倍此时它的尺寸与 skip 一致AttentionGate 接收(skip, x)把解码特征 x 作为门控信号 gskip 作为被加权的特征。拼接后接 DoubleConv完成这一尺度的特征融合。参数说明in_ch 是解码输入通道数skip_ch 是编码同层通道数先把 in_ch 减半再上采样是 U-Net 的经典通道折叠写法避免拼接后通道数直接翻倍导致显存失控。完整网络的 forward 决定了 AG 加在哪几层。以下结构里最深的两个尺度开启注意力浅的两层关闭因为 x1、x2 分辨率高、边缘信息重门控在这里容易把细小结构压掉class AttentionUNet(nn.Module): def __init__(self, in_ch1, out_ch1, base_ch64): super().__init__() self.inc DoubleConv(in_ch, base_ch) # 256 self.down1 Down(base_ch, base_ch * 2) # 128 self.down2 Down(base_ch * 2, base_ch * 4) # 64 self.down3 Down(base_ch * 4, base_ch * 8) # 32 self.down4 Down(base_ch * 8, base_ch * 8) # 16 self.up1 Up(base_ch * 8, base_ch * 8, base_ch * 4, use_attTrue) self.up2 Up(base_ch * 4, base_ch * 4, base_ch * 2, use_attTrue) self.up3 Up(base_ch * 2, base_ch * 2, base_ch, use_attFalse) self.up4 Up(base_ch, base_ch, base_ch, use_attFalse) self.outc nn.Conv2d(base_ch, out_ch, kernel_size1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) y self.up1(x5, x4) y self.up2(y, x3) y self.up3(y, x2) y self.up4(y, x1) return self.outc(y)结构说明最深层 down4 之后不再下采样x5 是 16×16×512 的语义特征。up1 和 up2 接编码器的 x4、x3这两层分辨率低、语义相对明确加 AG 收益明显up3、up4 接 x2、x1关闭门控以保留边界细节。outc 用 1×1 卷积把通道压到 out_ch配合 sigmoid 得到逐像素概率。3.3 最小训练循环与超参表训练部分默认用 BCE Dice 组合损失。一个最小 train step 如下def train_step(batch, model, optimizer, loss_fn): img, mask batch # (B, 1, 256, 256), (B, 1, 256, 256) img, mask img.cuda(), mask.cuda() optimizer.zero_grad() logits model(img) loss loss_fn(logits, mask) # 输出未过sigmoid损失内部处理 loss.backward() optimizer.step() return loss.item()逻辑说明logits 直接进损失函数而不是先做 sigmoid是为了避免 sigmoid 之后再取 log 造成数值不稳定BCE 分支应该在内部用 logits 计算Dice 分支再把 logits 过 sigmoid 得到概率。返回值是标量 loss便于每个 epoch 打印和早停判断。这套超参在多数 2D 分割任务上能直接起步超参数推荐值说明输入尺寸256×256原图过大时先裁剪到 ROI 附近不要全局 resizebatch size816 GB 显存的安全值OOM 时降到 4初始学习率1e-4Adam 下比 1e-3 稳后续配 cosine 衰减训练轮数100~200医学数据量小配合早停观察验证 Dice损失函数BCE Dice类别不平衡时的默认组合权重初始化KaimingPyTorch 默认遇 NAN 先查 BN 和归一化注意如果单卡只有 8 GB优先把输入降到 192 或缩短编码路径到 3 次下采样不要先动通道数。通道数减半会同时削弱编码和门控的表达力输入尺寸小幅下降只影响感受野上限通常损失更小。3.4 显存墙3D数据与2.5D折中CT 数据天然是三维体直接上 3D U-Net 要把所有 Conv2d 换成 Conv3d中间层体积是 C×D×H×W256 立方输入在中层就有几十亿浮点数单卡基本爆显存。常见做法是 2.5D取病灶中心附近的 3 张连续切片作为三通道输入或随机采样若干层做 2D 训练。在这套系统里先跑通 2D 再决定升级 3D。2D 训练稳定、可视化方便、迭代快3D 的收益主要体现在病灶跨层连续性的保持上。但如果数据层厚不均比如 5mm 层厚扫描z 方向的信息本身是错位的强上 3D 反而把噪声引进来这时候 2.5D 的三切片方案比纯 3D 更稳。4. 医学图像分割系统的预处理、损失函数与评价指标选型4.1 NIfTI与DICOM的读取与归一化医学数据最常见 NIfTI.nii.gz和 DICOM 目录两种来源。.nii 用 nibabel 读取DICOM 序列用 pydicom 或 SimpleITK 读取。读取后第一件事不是归一化而是看数值范围CT 的像素值是亨氏单位HU范围可达上千MRI 是相对信号强度没有固定单位。两者都不应该直接做 min-max要按模态语义处理。CT 一般做窗宽窗位截断把软组织映射到 0~1import numpy as np import nibabel as nib def load_nifti_with_window(path, ww400, wl40): img nib.load(path).get_fdata().astype(np.float32) lo wl - ww / 2.0 hi wl ww / 2.0 img np.clip(img, lo, hi) img (img - lo) / (hi - lo) return img逻辑说明ww 是窗宽决定显示的灰度范围wl 是窗位决定范围中心。腹部 CT 常用窗宽 400、窗位 40 突出软组织肺窗则用窗宽 1500、窗位 -600。clip 之后线性映射到 0~1既压制骨骼和空气的极端值又给网络一个稳定的输入分布。MRI 没有 HU 概念一般做均值为 0、标准差为 1 的 z-score或按 1% 和 99.5% 分位数截断。注意同一批数据必须用同一套归一化参数不能逐张算 min-max否则训练集和推理集的分布不一致。4.2 数据增强与样本划分医学分割的两个隐藏约束医学数据往往只有几十到几百例增强不是可选项。但从自然图像搬来的增强要改参数不要用旋转 90 度CT 扫描的上下语义固定不要大幅度随机缩放解剖结构的相对尺寸有临床意义。常用组合是小角度旋转、水平翻转、弹性形变alpha 控制在 5 以内import albumentations as A train_aug A.Compose([ A.RandomRotate90(p0.5), A.HorizontalFlip(p0.5), A.ElasticTransform(alpha3, sigma50, p0.3), A.RandomGamma(gamma_limit(80, 120), p0.2), ]) def apply_aug(image, mask): out train_aug(imageimage, maskmask) return out[image], out[mask]逻辑说明ElasticTransform 对每个像素施加随机位移场模拟组织形变是医学增强里最常用的非线性变换alpha 控制位移幅度3 已经偏保守超过 10 会把器官形状破坏掉。RandomGamma 模拟扫描仪亮度差异对 MRI 尤其有用。参数说明mask 与 image 必须走同一套变换albumentations 的 mask 参数会自动对齐手动实现时最容易出错的就是图和标注变换不一致。样本划分是另一个隐藏约束同一个病人的多层切片必须划到同一个集合里。CT 一个病例有几百张切片如果随机分训练验证模型等于见过同一病人的相邻切片验证集会虚高 3~5 个点。按 patient 维度分组切分是医学图像分割系统里最容易被忽视的基准线问题。4.3 损失函数前景占比不足5%时的选择逻辑分割任务的默认损失是交叉熵但医学数据的前景背景比例经常到 1:50CE 会被背景主导。选型优先级应该是先解决不平衡再解决边界模糊最后才考虑榜单排名。下表是几种常用损失的核心行为损失函数核心思路适合场景主要问题Dice Loss1 - 2TP/(2TPFPFN)类别不平衡前景小极端不平衡时梯度波动大BCE Dice两者直接相加默认组合最稳需要匹配权重Tversky对 FP/FN 分别加权优先降漏检FN多一个要调的 betaFocal LossCE × (1-p)^γ难样本比例高γ 调不好容易掉点Boundary Loss距离图上的积分边界质量优先实现复杂数值不稳选择逻辑第一版直接上 BCE Dice。Dice 对前景大小不敏感BCE 提供平滑梯度两路相加通常比单独 Dice 收敛快且稳。如果临床反馈漏检多把 Tversky 的 beta 设到 0.7朝 FN 方向加权如果边界毛糙再考虑 Boundary Loss但不要第一版就上。每个损失都要配合验证集 Dice 看趋势而不是只看训练 loss 下降。4.4 Dice、IoU与HD95的评估代码训练中要盯的指标有三个Dice、IoU 反映重合度HD95 反映边界偏差。前两个直接实现import numpy as np def dice_coef(pred, gt, eps1e-6): inter (pred * gt).sum() return (2 * inter eps) / (pred.sum() gt.sum() eps) def iou_coef(pred, gt, eps1e-6): inter (pred * gt).sum() union pred.sum() gt.sum() - inter return (inter eps) / (union eps)逻辑说明pred 和 gt 在传入前先转成 0/1 二值数组通常取 pred 0.5 做阈值。eps 防止分母为 0但对稀疏前景来说eps 取 1e-6 足够。HD95 计算两个集合中每个边界点到另一个集合的最短距离并取 95 分位标准做法是直接用 medpy.metric.binary.hd95传入两个二值数组即可。注意3D 体积评估时不要按 2D 切片先求 Dice 再平均。切片级平均和体积级 Dice 的数值在内部监控里差别不大但论文对比中会被审稿人质疑。统一口径3D 数据就在 3D 体积上算一次2D 数据就在 2D 上算一次。5. Attention U-Net的注意力系数提取与分割结果验证技巧5.1 用forward hook把注意力图抽出来验证 Attention U-Net 是否学到东西直观的证据是注意力图本身。不需要改动训练代码注册 forward hook 即可alpha_maps {} def hook_fn(name): def fn(module, inp, out): alpha_maps[name] out[1].detach().cpu() # (B, 1, H, W) return fn ag_layers [m for m in model.modules() if isinstance(m, AttentionGate)] for i, ag in enumerate(ag_layers): ag.register_forward_hook(hook_fn(fag_{i}))逻辑说明AttentionGate.forward 返回(alpha * x, alpha)out[1] 正好是注意力系数。hook 在每层 AG 前向结束后把 alpha 存进 dict推理结束后取出即可叠加可视化。参数说明name 用 AG 在模型中的顺序索引区分B 为 batch sizeH、W 是当前尺度的分辨率浅层 AG 的 H、W 大深层小可视化前要上采样到同一尺寸。5.2 注意力图的三种典型失效模式拿到 alpha 后把它双线性上采样到原图尺寸用半透明 colormap 叠加到原始影像上。常见的三种现象对应三个不同问题alpha 整体接近 1没有空间选择性。门控没训练起来最常见原因是学习率偏大导致 self.wx、self.wg 的梯度爆炸或者轮数太少。先降到 5e-5 重训不要动结构。alpha 在背景区域也有大范围高响应。说明背景纹理和目标特征严重相似。优先检查预处理归一化范围是否统一、增强是否引入了不自然的纹理而不是换损失。alpha 与标注位置吻合但比 mask 大一圈。网络学到的是目标邻域而不是目标本身常见于边界模糊的病灶。这时换带 FN 权重的 Tversky 比换网络更有效。这步排查的价值在于把黑盒分割变成一个可解释的决策过程。网络在哪一层、以哪种方式犯错从 alpha 的空间分布能直接看出来比盯着 Dice 数值猜原因高效得多。5.3 一个固定的验证流程最后给一套我在项目里固定使用的验证顺序。先固定随机种子训练一份模型保存验证集每个样本的 Dice、HD95 和全部层级的 alpha。然后做三件事第一统计每个 AG 层 alpha 的均值与方差训练收敛后均值应该在 0.3~0.8 之间如果全部接近 0说明门控把目标区域整体抑制了属于败坏的训练第二随机挑 20 个样本把 alpha 热力图和预测 mask 并排导出人工确认高注意力区域与分割前景一致第三做 5 次不同种子的重复实验对比加 AG 与不加 AG 的 Dice 差异差异小于 0.01 时优先选结构更简单、推理更快的普通 U-Net。Attention U-Net 的收益集中在低对比度、小目标场景模型结构的价值要放在具体数据分布里验证而不是靠注意力图好看程度来评价。本文还有配套的精品资源点击获取