简介这是一套面向医学影像分割方向研究者的改进型U-Net实现方案聚焦人体脊椎分割任务在经典U-Net中引入SE通道注意力模块与Transformer自注意力机制用于提升对复杂脊椎形态和边缘细节的分割精度。压缩包共2000个文件其中1501个png和492个jpg构成训练与测试图像数据集4个py文件为核心模型与训练脚本另有项目说明书docx、环境说明txt等文档包体仅26.89MB便于快速部署和复现。目前已有73人学习下载。资源附带完整的项目说明书清晰说明改进原理、数据标注规范与使用流程数据集覆盖多种脊椎影像样例可用于模型训练与验证。对正在研究CNN与Transformer结合、或需要高质量脊椎分割代码与标注数据的开发者而言这套资料能直接提供可运行的工程基础省去从零搭建和整理数据的时间。1. 基于U-Net改进SE和Transformer的人体脊椎分割这套组合拳到底解决了什么做过医学影像分割的人多半有过这种体验拿一份标注好的椎骨CT数据用经典U-Net训练Dice能到0.85以上但把模型部署到另一个扫描协议的数据上椎体边缘就开始出现锯齿状误检椎间盘和椎体边界糊成一团。这个标题里的思路——在U-Net骨架里插入SE通道注意力机制再在瓶颈层接上Transformer模块——正是冲着这类问题去的SE负责让模型关注“哪些特征通道该被放大”Transformer负责弥补卷积局部感受野带来的上下文缺失让模型在分割椎骨、椎间盘、棘突等解剖结构时不再只盯局部纹理。适合正在做脊柱影像AI项目、论文改进实验或医学影像竞赛的人它给的不是全新网络而是在成熟U-Net框架上可复现的模块化改造路线。2. 先把U-Net跑通再做改进基线训练与解剖结构约束在做任何模块改进之前我强烈建议你先在原版U-Net上跑出一个能用的基线再把SE和Transformer加进去。没有基线后面所有消融实验都是空中楼阁审稿人或者你的leader问“改进比基线高多少”时你拿不出数字整个项目就失去了说服力。2.1 为什么脊柱分割默认选U-Net而不是别的分割网络人体脊椎分割的数据形态和自然图像分割有明显差异。CT影像里椎体、椎间盘、棘突的灰度值彼此接近边缘不像自然图像那样有清晰的色彩跳跃而是靠骨窗和软组织窗的窗宽窗位来凸显。U-Net在这个场景下依然是首选原因有三第一它通过跳跃连接保留了高分辨率特征图椎体这类小结构边缘不会在逐层下采样中丢失第二编码器-解码器结构天然适配医学影像这种“输入输出尺寸一致”的需求不需要额外的反卷积头第三U-Net的参数量适中在单张12G显卡上就能训练一个不错的2D模型这对很多没有多卡环境的实验室来说是硬约束。但原版U-Net有个明显短板编码器每一层都用卷积池化下采样感受野虽然在扩大但仍然是局部窗口的堆叠。椎体分割中最常遇到的“相邻椎体边界粘连”问题本质上就是因为模型在深层特征图上缺乏全局上下文推理能力不知道“这一片椎体的边缘应该在哪里收住”。这就是后面要引入Transformer的原因不过这一步先按下不表。2.2 跑通最小训练流程代码结构与参数设置先给出一个可以直接复用的原始U-Net训练入口。这里我以PyTorch为例假设你的数据集已经整理成了标准的“图像-标签”对图像是单通道的CT切片标签是单通道的掩码像素值为0、1、2代表背景、椎体、椎间盘或其他自定义类别。# train_baseline.py import torch import torch.nn as nn from torch.utils.data import DataLoader from monai.losses import DiceLoss # 以monai提供的UNet为例实际项目中常自行实现或改用segmentation_models_pytorch from monai.networks.nets import UNet model UNet( spatial_dims2, in_channels1, out_channels3, channels(16, 32, 64, 128, 256), strides(2, 2, 2, 2), num_res_units2, ) loss_function DiceLoss(to_onehot_yTrue, softmaxTrue) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max200) for epoch in range(200): model.train() for batch in train_loader: images, labels batch[image], batch[label] logits model(images) loss loss_function(logits, labels) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() if epoch % 10 0: print(fepoch {epoch}, loss {loss.item():.4f})这里有几个参数需要特别留意。channels从16起步逐层翻倍到256这是U-Net的常规设计但如果你训练的CT切片分辨率是512×51216个初始通道可能稍显不足可以把第一层通道数调成32或64代价是显存占用上升训练速度下降。num_res_units2表示每个编码/解码块内堆叠2个残差卷积单元这个值从0到2都有人用取值越大模型拟合能力越强但脊柱分割这种结构相对规则的任务2个单元已经足够再往上加反而容易在小数据集上过拟合。如果在这个基线阶段就发现Dice上不去不要在模型结构上找原因先检查你的数据预处理窗宽窗位是否统一、标签和图像是否对齐、切片的spacing是否一致。这类问题占脊柱分割项目初期翻车的八成原因。3. SE通道注意力加在哪从易混淆椎体看通道重标定的实际效果SE模块Squeeze-and-Excitation在2017年被提出后很长一段时间内都是CNN改进的标配。它的核心逻辑很朴素卷积层输出的每个特征通道重要性不一样有的通道编码的是骨边缘有的通道编码的是软组织纹理SE通过学习每个通道的权重让模型把注意力集中在更有判别力的特征上。放在U-Net里SE要解决的问题就是“椎体和椎间盘在灰度上太像了普通卷积不知道该放大哪些特征差异”。3.1 SE模块的结构和计算代价一个标准的SE模块包含两部分Squeeze操作通过对每个通道做全局平均池化把H×W的空间信息压缩成一个标量得到C维向量Excitation操作通过两层全连接第一层降维到C/r第二层升维回C和Sigmoid激活生成每个通道的权重然后把这个权重乘回原始特征图的每个通道。这里的r叫缩减率常见取值是16含义是中间隐藏层维度压缩到C/16减少参数量。# se_block.py import torch import torch.nn as nn class SEBlock(nn.Module): def __init__(self, channels, reduction16): super().__init__() self.squeeze nn.AdaptiveAvgPool2d(1) self.excitation nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels), nn.Sigmoid(), ) def forward(self, x): b, c, h, w x.shape y self.squeeze(x).view(b, c) y self.excitation(y).view(b, c, 1, 1) return x * y这段代码是实现SE最直接的写法。AdaptiveAvgPool2d(1)把每个通道的特征图压成一个数感受野覆盖整个特征图这一步提供的是全局统计信息。ReLU和Sigmoid的位置不能换前者负责引入非线性并限制中间层表达后者把输出压缩到0到1之间表示每个通道的“保留比例”。计算量上SE只增加少量全连接层参数对整体推理速度的影响几乎可以忽略这也是它能成为“免费午餐”式改进的原因。3.2 SE放在编码器末端还是跳跃连接处我的推荐路径直接说结论在脊柱分割这个任务上我建议跳过把SE放在跳跃连接上的流行做法把它放在编码器最后一个block的输出上、以及解码器每个block的输入之前这样效果最稳定。原因在于跳跃连接传递的是编码器高分辨率特征这些特征里既有细节纹理也有噪声脊柱CT里常见的软组织伪影会混在其中。如果在这里加SE模块会把通道权重聚焦到当前切片最显著的纹理上——而这往往是伪影最重的区域结果是你过滤了有用的边缘特征。网上很多教程把SE放在跳跃连接之后但这个做法在脊柱数据上我实测是掉点的属于典型的“看着合理实则翻车”的改进。放在编码器末端和解码器输入处就不同了。编码器末端特征图分辨率为16×16或8×8通道数最多语义信息最丰富SE在这里进行通道重标定等于在告诉模型“深层特征里边缘特征通道要比背景特征通道更值得放大”。具体代码接入方式如下# unet_with_se.py核心片段 class EncoderBlockWithSE(nn.Module): def __init__(self, in_ch, out_ch, se_reduction16): super().__init__() self.conv1 nn.Conv2d(in_ch, out_ch, 3, padding1) self.norm1 nn.BatchNorm2d(out_ch) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(out_ch, out_ch, 3, padding1) self.norm2 nn.BatchNorm2d(out_ch) self.se SEBlock(out_ch, se_reduction) def forward(self, x): x self.relu(self.norm1(self.conv1(x))) x self.relu(self.norm2(self.conv2(x))) x self.se(x) return x这里有几个关键点值得展开。第一SE必须在第二个卷积和归一化之后接入不能放在第一个卷积后否则注意力权重会被浅层特征主导起不到筛选深层语义的作用。第二se_reduction这个参数对脊柱分割影响很大设置过小比如4SE的参数量暴涨在小样本数据上很容易过拟合设置过大比如32通道间的差异性被过度压缩权重趋同SE退化成恒等映射。我在腰椎数据集上的经验是reduction16是稳的如果你的类别数多、通道数大可以尝试8。第三SE输出之后不需要额外接残差连接因为SEBlock内部没有改变特征图的形态权重乘回去就是增强后的特征。4. Transformer模块接在瓶颈层还是跳跃连接显存与感受野的权衡Transformer这块是这个项目里最容易被做成“花架子”的部分。很多人一听到Transformer就觉得要把整个U-Net换成Swin-Unet但摊开看脊柱分割的数据量级往往在几百到几千张切片在这种规模下上全局Transformer收益有限显存却先爆了。正确的做法是把Transformer当成一个模块插到U-Net最需要全局上下文的位置。4.1 瓶颈层接Transformer16×16分辨率是安全区U-Net的瓶颈层是特征图分辨率最低、通道数最高的位置原始尺寸512×512的输入下经过四次下采样后变成32×32甚至16×16的标准特征图这个分辨率对Transformer来说是友好的。Transformer的核心计算量在Self-Attention上复杂度是序列长度的平方如果直接在32×32的特征图上做全局注意力序列长度是1024注意力矩阵是1024×1024再加多头训练时显存会瞬间吃掉好几G。16×16分辨率下序列长度只有256注意力矩阵小了一个量级这才是视觉Transformer在医学分割里能落地的原因。在瓶颈层接Transformer我一般选择两个方案之一。方案一是直接用标准Transformer Encoder Layer把CNN输出的特征图按通道展开成序列送入Multi-Head Attention和Feed-Forward Network再重塑回特征图方案二是走Swin Transformer路线用窗口划分降低注意力复杂度但窗口尺寸、移位次数的超参数调起来比较麻烦项目周期紧时不建议在U-Net的瓶颈层引入这种复杂度。下面给出方案一的参考实现# bottleneck_transformer.py import torch import torch.nn as nn class TransformerBottleneck(nn.Module): def __init__(self, dim, num_heads8, mlp_ratio4., dropout0.1): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention(embed_dimdim, num_headsnum_heads, dropoutdropout, batch_firstTrue) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(dropout), ) def forward(self, x): b, c, h, w x.shape x x.flatten(2).transpose(1, 2) # (b, h*w, c) x x self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x x self.mlp(self.norm2(x)) x x.transpose(1, 2).view(b, c, h, w) return x写这段代码有三个容易踩的坑。第一个坑是nn.MultiheadAttention默认的batch_firstFalse如果你不熟悉这个API很容易把输入维度搞反报错时看到“expected shape”一脸懵最好的习惯是显式设置batch_firstTrue并在forward里保证输入是(batch, seq_len, dim)。第二个坑是LayerNorm必须在Attention之前做这是Pre-Norm结构和最早的Post-Norm Transformer不同——Pre-Norm训练更稳定在医学影像小数据集上几乎不会出现训练崩溃。第三个坑是mlp_ratio4也就是FFN隐藏层是输入维度的4倍这是Transformer系列的默认配置改大改小都会影响表达能力和训练速度。4.2 把Transformer用在跳跃连接上的代价有一部分改进思路是把Transformer接在解码器的跳跃连接上让解码器每层都能获得全局上下文信息。这个思路在理论上是成立的但实践中的代价极大。以512×512输入、四层U-Net为例第一层跳跃连接的特征图分辨率是256×256序列长度是65536这个数量级下任何全局注意力都会直接把显卡内存吃光即使切成局部窗口窗口之间的交互也会让实现复杂度上升好几个台阶。我不是说这个方向没法做而是在“基于U-Net改进SE和Transformer”这个项目定位下它超出了性价比边界。我的建议是要加Transformer就只加在瓶颈层最多加到倒数第二层下采样后的特征图上分辨率32×32时用窗口注意力。跳跃连接处继续保留原有的拼接操作让SE去处理通道维度的筛选这样整个模型的改进点是清晰且可控的。当你做消融实验时可以说“Transformer模块引入了全局上下文建模SE模块增强了通道判别力二者作用在不同维度互补冲突小”这是审稿人或者答辩老师最想听到的话。5. 逐层改进的完整训练管线把SE和Transformer拼进一个模型讲完两个模块各自的位置这一章把它们拼成一个完整的可训练模型。在结构设计上我采用的是一个典型的“编码器-SE增强-Transformer瓶颈-解码器-SE增强”的U形结构建模思路是保留U-Net的骨架不变在关键位置插入两个模块这样既能兼容已有的U-Net训练代码又能让你方便地做消融实验。5.1 最终网络结构定义与完整前向逻辑下面这个模型定义整合了前面两章的SEBlock和TransformerBottleneck编码器和解码器分别使用残差卷积块并在编码器末端、解码器输入端插入SE在瓶颈层插入Transformer。为了保持代码简洁这里用nn.ModuleList管理不同分辨率的编码器块。# spu_net.py import torch import torch.nn as nn from se_block import SEBlock from bottleneck_transformer import TransformerBottleneck class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch, use_seFalse): super().__init__() self.conv1 nn.Conv2d(in_ch, out_ch, 3, padding1) self.norm1 nn.BatchNorm2d(out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, padding1) self.norm2 nn.BatchNorm2d(out_ch) self.relu nn.ReLU(inplaceTrue) self.se SEBlock(out_ch) if use_se else nn.Identity() def forward(self, x): x self.relu(self.norm1(self.conv1(x))) x self.relu(self.norm2(self.conv2(x))) x self.se(x) return x class SPUNet(nn.Module): def __init__(self, in_channels1, out_channels3, base_channels32, use_seTrue, use_transformerTrue): super().__init__() # 编码器 self.enc1 ConvBlock(in_channels, base_channels) self.enc2 ConvBlock(base_channels, base_channels * 2) self.enc3 ConvBlock(base_channels * 2, base_channels * 4) self.enc4 ConvBlock(base_channels * 4, base_channels * 8, use_seuse_se) self.pool nn.MaxPool2d(2) # 瓶颈SE Transformer self.bottleneck nn.Sequential( ConvBlock(base_channels * 8, base_channels * 16, use_seuse_se), TransformerBottleneck(dimbase_channels * 16, num_heads8) if use_transformer else nn.Identity(), ) # 解码器 self.up4 nn.ConvTranspose2d(base_channels * 16, base_channels * 8, 2, stride2) self.dec4 ConvBlock(base_channels * 16, base_channels * 8, use_seuse_se) self.up3 nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride2) self.dec3 ConvBlock(base_channels * 8, base_channels * 4, use_seuse_se) self.up2 nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride2) self.dec2 ConvBlock(base_channels * 4, base_channels * 2, use_seuse_se) self.up1 nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride2) self.dec1 ConvBlock(base_channels * 2, base_channels, use_seuse_se) self.out_conv nn.Conv2d(base_channels, out_channels, 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)) bn self.bottleneck(self.pool(e4)) d4 self.up4(bn) d4 torch.cat([d4, e4], dim1) d4 self.dec4(d4) d3 self.up3(d4) d3 torch.cat([d3, e3], dim1) d3 self.dec3(d3) d2 self.up2(d3) d2 torch.cat([d2, e2], dim1) d2 self.dec2(d2) d1 self.up1(d2) d1 torch.cat([d1, e1], dim1) d1 self.dec1(d1) return self.out_conv(d1)这段代码把结构串联得很清楚重点解读三个设计选择。第一是SE只在enc4、bottleneck和所有解码器块中启用enc1到enc3不加SE。因为低层特征图分辨率高、通道少加SE的收益不明显反而增加过拟合风险你可以把这个设计当作一个可调项如果发现模型在验证集上高分辨率特征贡献不足再把SE往前移一层试一次。第二是瓶颈层先接一个ConvBlock再接Transformer这个顺序是故意的——卷积先完成局部特征的点化Transformer在整个特征图的所有位置上做上下文交互两个模块各司其职而不是功能重叠。第三所有跳跃连接依然保留原始的concat操作这是U-Net的核心优势不要因为加了Transformer就丢掉跳跃连接。5.2 损失函数、训练策略与数据划分网络结构定了训练策略也要跟着调整。脊柱分割最常用的损失函数是DiceLoss和CrossEntropyLoss的组合这是因为DiceLoss直接优化系数本身而CE能为每个像素提供梯度两者结合可以减少训练初期的收敛震荡。下面的代码给出一个典型的混合损失用法# loss.py import torch import torch.nn as nn import torch.nn.functional as F from monai.losses import DiceLoss class DiceCEWeightedLoss(nn.Module): def __init__(self, alpha0.7): super().__init__() self.dice DiceLoss(to_onehot_yTrue, softmaxTrue) self.ce nn.CrossEntropyLoss() self.alpha alpha def forward(self, logits, labels): # logits: (b, c, h, w), labels: (b, h, w) or one-hot dice_loss self.dice(logits, labels) ce_loss self.ce(logits, labels.long()) return self.alpha * dice_loss (1 - self.alpha) * ce_loss这个混合损失里alpha0.7是我在腰椎CT数据集上调出来偏好Dice占比更大更适配“前景占比小、背景占比大”的医学影像特点。如果你发现椎间盘这类小目标被整体漏掉了可以考虑加一个类别的加权项在CE里给椎间盘类别更高权重——常用的做法是算训练集各类别像素占比的倒数归一化。训练轮数上加入Transformer后收敛速度比纯CNN慢原版U-Net可能80个epoch就收敛了改进版可能要150到200个epoch这是正常的不要中途停下来说模型坏了。数据划分也有一个血泪经验必须按“病人”而不是按“切片”划分训练集和验证集。脊柱CT一个病人有几十张到几百张切片如果随机按切片划分同一个病人的相邻切片被分到训练集和验证集验证指标会虚高得像模型已经见过答案测试和部署时立刻原形毕露。按病人维度划分后训练集和验证集的分布差异会更大这也才是改进模块真正能体现价值的场景。5.3 训练启动前的检查清单启动训练前建议跑一次单batch的过拟合测试。方法很简单取训练集一个batch把模型迭代几百步看训练损失能不能降到接近0。如果降不下去先别加Transformer退回纯U-Net基线排查因为问题可能出在标签对齐或者归一化上而不是新引入的模块。这一步能省掉你后面排查的一堆脏数据和错误属于花十分钟换两小时的经验。6. 验证技巧与改进兜底拿这两个指标判断你的改进是否翻车项目说明书里最重要的部分除了结构和代码就是验证方案。模型改进得再花哨最终要靠指标说话这也是区分“项目实战”和“玩具demo”的分水岭。这一章我讲两个核心指标的验证方法以及SE和Transformer真正在脊柱分割上值得关注的边界。6.1 不只盯Dice边界距离指标HD95才是硬标准Dice系数是区域重叠度的度量它有一个众所周知的盲区当两个分割结果的重叠区域面积相近但边界位置不同时Dice可能差异不大但实际解剖位置偏差已经很大。对脊柱分割来说椎体边界偏移几个像素就可能导致手术导航定位不准这种误差用Dice根本看不出来。所以在项目说明书的验证部分除了Dice几乎一定会包含HD95——95%分位数Hausdorff距离它表示预测边界和真实边界之间绝大多数对应点距离都小于某个阈值。代码层面可以直接用monai的HausdorffDistanceMetric模块算# evaluate.py from monai.metrics import HausdorffDistanceMetric import numpy as np hd_metric HausdorffDistanceMetric(include_backgroundFalse, percentile95) # 假设 pred 是模型输出的argmax结果shape(b, h, w) # target 是one-hot形式的标签 hd_values hd_metric(y_predpred, ytarget) print(HD95:, hd_values.mean().item())跑这个指标时要注意percentile95对应的就是HD95如果项目说明书写得很粗糙没有明确百分位默认可能算的是最大Hausdorff距离那个值对边界噪声极其敏感一个离群点就能把指标拉爆不能拿来作为稳定报告。另一个坑是include_backgroundFalse这个参数含义是排除背景类别的距离计算只计算前景椎体和椎间盘的距离如果你漏了这个参数背景的占比会让数值虚低误导结论。我的习惯是同时报告Dice和HD95Dice提升、HD95下降说明改进是实打实的如果Dice不变HD95降了说明边缘更锐利也是有效的如果两个指标都变差那就得回到结构设计去找问题了。6.2 消融实验的对照表判断SE和Transformer各自的净收益项目说明书里通常要求提供消融实验表这是整个项目里投入产出比最高的三行。设计三组对照原版U-Net、U-NetSE、U-NetSETransformer。这样能明确回答两个问题——“SE改了通道注意力机制带来的收益有多少”和“Transformer在SE基础上还能往上提多少”。我在实际项目里常见的结果是SE单独带来Dice提升约0.5到1个百分点而Transformer叠加之后的提升在1到2个百分点。如果SE的提升是负数需要回头检查加SE的位置是否选在了跳跃连接或低层特征上如果Transformer叠加后的提升不明显先确认你训练的轮数够不够Transformer的收益经常要到200个epoch之后才显现这属于它的建模特性不是模型坏了。还有一个在验证阶段容易被忽视的细节类别级别的指标拆分。脊柱分割的标签里通常椎体和椎间盘被标成不同类别把Dice拆到类别层面你才能看到Transformer对上下文建模的收益到底落在流体结构上。椎间盘的Dice提升通常是Transformer模块最明显的正反馈因为它的形态和相邻椎体的位置关系高度依赖上下文推断——只靠CNN的局部特征很难把“两个椎体中间的那条窄缝”分对而Transformer的全局注意力正好补上了这一块。如果做完类别拆分后发现椎间盘Dice没涨反而有下滑我通常会检查数据增强里是否加入了过强的随机弹性形变因为Transformer对空间位置的敏感性比CNN高过度的几何形变会让它学到的位置编码失效。收尾想分享一个自己的习惯每次做通道注意力和Transformer的改进前我都要先在原版模型上跑一遍完整训练和验证记录下每个类别的Dice和HD95把这个结果存成“底线文件”。改完之后翻车了就回去对照底线看模块加在哪一层影响最大这个对照动作帮我少交了不少“看起来高大上但实际不如原版”的学费。每回有人问起这类U-Net改进值不值得做我都会说值得但前提是先把基线做实、把验证做全否则注意力机制和Transformer带来的不是提升而是一堆解释不清的参数和没完没了的调参。这套做法放到你的脊柱分割项目里同样适用希望帮到你。本文还有配套的精品资源点击获取