医学图像分割这一块UNet几乎是默认起点。我用UNet做过不少任务从CT里的器官分割到显微镜下的细胞分割多数时候表现不错但总有一些情况让人挠头小目标容易丢、边界模糊的地方分不清、深层的语义特征和浅层的细节特征像是两张皮。这也是我一直坚持做UNet改进的原因。今天这篇是系列第43篇把坐标注意力Coordinate Attention接进UNet专门针对医学图像分割再做一轮优化。这篇文章不是原理复读我会从“为什么之前的注意力不够用”讲起拆坐标注意力的内部流程再把模块放进UNet的实操细节和调参经验一起整理出来适合正在改分割模型、特别是跟医学影像里的细线结构和边界模糊较劲的读者。1. 为什么UNet在医学图像分割里“差一口气”1.1 医学图像分割的核心困难医学图像分割和自然图像分割最大的区别在于目标本身不“显眼”。自然图像里的物体通常有清晰的边缘和颜色差异而医学影像大多是灰阶图器官和病灶在灰度上和周围组织非常接近。以CT为例某些软组织的边界在肉眼看来都只是一层淡淡的灰度过渡超声图像里噪声和阴影会进一步掩盖目标边界血管造影里许多细分支只有两三个像素宽。这些问题决定了模型不能只靠局部纹理做判断还得有全局位置感知能力。另一个困难是目标尺寸跨度大。同一个模型既要能分割一个大的器官轮廓又要能捕捉一个几毫米的早期病变。UNet编码器逐层下采样后高层的感受野变大了但空间分辨率低了小目标细节很容易被“平均”掉。跳连接可以把浅层高分辨率特征送到解码器但如果深浅层特征在语义上没有对齐直接拼接反而会把噪声带进去让边界预测变得更碎。1.2 UNet的结构优势与明显短板UNet的优势不用多说编码器逐层提取语义解码器逐层恢复细节跳连接让梯度回传更顺畅。但连续的卷积本质上是局部算子每个位置的输出主要受附近一小块区域影响。即便网络足够深下采样在一定程度上扩大了有效感受野但不同空间位置之间的长程依赖仍然偏弱。打个比方可以把编码器特征看作一张城市地图卷积操作是在每个街区看门牌层数越深能看到的地图越广但缺少一个全局动作告诉模型“关键地标在哪个方向”。SE通道注意力能给每个通道算一个全局总分却把位置信息抹掉了CBAM虽然有空间注意力但它的空间权重对所有通道共享而且卷积核感受野有限。我先后试过这两种方案提升有但总觉得差口气尤其是对那些形状细长、分布在图像边缘的目标。1.3 坐标注意力为什么能对症坐标注意力最吸引我的一点是它把“位置”以很轻量的方式揉进了通道注意力。它不直接生成二维空间权重而是沿高度和宽度方向分别做池化得到两个带坐标的一维向量再组合起来。这样一来模型在增强“哪些通道重要”的同时也能知道“这些通道重要的大概位置在哪一行、哪一列”。对医学图像分割来说这个特性直接命中了痛点。目标经常偏在图的一侧病灶常常呈细长形态位置信息一旦丢失分割结果就会在远离目标的地方产生假阳性或者在细小分支处断掉。坐标注意力让网络多了一个位置感知维度而且模块是即插即用的输入输出形状不变、通道数不变放在UNet的任意卷积块后面都不需要改结构。2. 坐标注意力先看原理再谈改进2.1 从SE到坐标注意力到底改了什么东西先回忆SE模块对输入特征做全局平均池化得到每个通道的一个标量再用两个全连接层得到通道权重乘回原特征。整个过程就是“每个通道算一个总分所有空间位置共享这个权重”。问题在于一个通道在不同位置可能代表完全不同的语义。如果病灶只出现在图像的右下角SE会把右下角的响应和整张图的均值混在一起空间细节就这样被压掉了。CBAM在SE基础上加了空间注意力分支把通道均值池化和最大池化结果拼在一起过一个卷积层生成空间权重。空间信息保住了但这个权重对所有通道是共享的不能体现“不同通道关注不同位置”的差异。那个卷积核的感受野也很有限对小目标依旧不友好。坐标注意力的做法是在SE的“通道总分”基础上增加了坐标标签不是压缩成一个标量而是在高度方向压缩宽度在宽度方向压缩高度得到两个带方向信息的一维向量。随后把两个方向拼接起来学习交互关系再分别生成注意力权重。这样既能保留通道重要性又能让模型知道目标更偏向哪一行哪一列。2.2 模块内部的完整流程落地时整个流程可以拆成五步。第一步输入特征图尺寸为B×C×H×W分别对高度方向和宽度方向求平均。对高度方向求平均得到形状B×C×H×1的张量对宽度方向求平均得到B×C×1×W的张量。注意这里没有像SE一样把空间压成一个点而是只压缩了一个方向。第二步把宽度方向的张量转置成B×C×W×1再和高度方向的张量在长度维度上拼接得到B×C×(HW)×1。这一步的意义是让水平信息和垂直信息进入同一个卷积层从而学习两者的交互关系。第三步用一个1×1卷积把通道数降到中间维度降低参数量再经过BatchNorm和ReLU激活。中间维度一般取输入通道除以reduction同时设一个下限比如8避免通道太少变成信息瓶颈。第四步将降维后的特征重新拆成两个张量分别恢复原来的通道数然后用Sigmoid生成水平方向和垂直方向的注意力权重。第五步把原始特征图分别乘上这两个方向的权重。由于两个权重的形状分别是B×C×H×1和B×C×1×WPyTorch广播机制会自动把它们扩展到B×C×H×W再逐元素相乘。这个流程里最关键的是“先拼接、再拆分”。正是这个设计让两个方向的信息先在中间特征里发生了交互而不是各算各的。如果去掉拼接直接独立生成方向权重模块会退化成一个很弱的变体效果会差不少。2.3 与SE、CBAM的轻量级对比我整理了一份从实际实验里得到的大致对比参数和计算量以基础UNet为参考。模块是否保留位置信息参数量增幅计算量增幅对小目标友好度SE否低很低中等CBAM是但空间权重共享中较高中等坐标注意力是且按通道区分低较低较高坐标注意力的参数主要集中在两个1×1卷积上reduction设为16或32时相对UNet的整体参数几乎可以忽略。CBAM的空间注意力部分虽然只有一个卷积层但会在训练时产生额外激活值显存上涨更明显。医学图像任务的数据量往往不大batch也小显存和参数量都得精打细算这是我后来选择坐标注意力的现实原因之一。3. 把坐标注意力“装”进UNet的正确姿势3.1 三个可插入位置实测结果差别很大拿到一个即插即用的模块最需要考虑的问题不是“怎么写”而是“放在哪”。我在UNet里认真试过三种位置。第一种是放在编码器每个stage的末端也就是两次卷积之后、下采样之前。这样可以保证下采样前特征已经被注意力增强降采样过程不会把关键响应过早弄丢。第二种是放在解码器每个stage的输出端也就是两次卷积之后再上采样之前。我实验下来这个位置比放在上采样之后更稳定因为上采样操作会改变空间尺寸先增强再放大细节恢复更可控。第三种是放在跳连接上。跳连接是UNet的命脉在这里加注意力可以压制浅层特征里与目标无关的噪声但代价是每一条跳连接都多一个模块显存上涨明显。实测下来对边界精细类任务有提升对常规大器官任务提升不大所以我默认只在需要精细边界的任务里打开这个开关。最后我的默认结构是编码器每个stage加CA解码器每个stage加CA跳连接不加。这个组合在多数医学分割任务上性价比最高。只加编码器也能跑但解码器部分在恢复被降采样抹掉的位置细节时有了CA后边缘明显更锐利。3.2 一个可以直接复用的CA模块代码下面这段PyTorch实现是我常用的版本。为了便于在UNet里替换我做了点简化没有引入太特别的激活函数输出通道和输入通道保持一致。import torch import torch.nn as nn class CoordinateAttention(nn.Module): def __init__(self, in_channels, reduction32): super().__init__() mip max(8, in_channels // reduction) self.conv1 nn.Conv2d(in_channels, mip, kernel_size1, biasFalse) self.bn1 nn.BatchNorm2d(mip) self.relu nn.ReLU(inplaceTrue) self.conv_h nn.Conv2d(mip, in_channels, kernel_size1, biasFalse) self.conv_w nn.Conv2d(mip, in_channels, kernel_size1, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): b, c, h, w x.shape x_h torch.mean(x, dim3, keepdimTrue) # b,c,h,1 x_w torch.mean(x, dim2, keepdimTrue).permute(0, 1, 3, 2) # b,c,w,1 y torch.cat([x_h, x_w], dim2) # b,c,hw,1 y self.relu(self.bn1(self.conv1(y))) x_h, x_w torch.split(y, [h, w], dim2) x_w x_w.permute(0, 1, 3, 2) # b,c,1,w a_h self.sigmoid(self.conv_h(x_h)) a_w self.sigmoid(self.conv_w(x_w)) return x * a_h * a_w代码已经跑通过。torch.mean是无参数的不会影响梯度计算。split时使用原始输入的H和W做长度切分所以输入尺寸不固定也没关系。3.3 在UNet里的组合示例为了方便说明我截取一个编码器stage的写法。假设输入64通道的特征图经过两个卷积块后接CA再下采样。class EncoderStage(nn.Module): def __init__(self, in_ch, out_ch, with_attTrue): super().__init__() self.conv1 nn.Conv2d(in_ch, out_ch, 3, padding1) self.bn1 nn.BatchNorm2d(out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, padding1) self.bn2 nn.BatchNorm2d(out_ch) self.relu nn.ReLU(inplaceTrue) self.att CoordinateAttention(out_ch) if with_att else nn.Identity() self.pool nn.MaxPool2d(2) def forward(self, x): x self.relu(self.bn1(self.conv1(x))) x self.relu(self.bn2(self.conv2(x))) x self.att(x) return x, self.pool(x)这里返回值有两个一个是送给跳跃连接的特征另一个是下采样后的特征。我习惯把CA放在两个卷积之后、下采样之前这样跳跃连接拿到的同样经过了注意力增强对解码器更友好。你可以按这个模式把CA插到UNet每个stage尾部不需要改动其他部分。4. 实验设置与效果数据4.1 数据和指标怎么定的因为涉及真实影像数据来源不方便公开。我用了两个内部整理的医学分割任务做验证。第一个是腹部CT影像里的目标结构分割共200例样本目标在图像中的占比变化很大典型的类不平衡情况。第二个是造影图像里的细线状结构分割共150例样本很多目标只有两三个像素宽非常考验位置感知能力。数据统一按8:2划分训练和验证验证集使用相同的数据增强流程来保证公平。评价指标方面我主要看Dice系数和IoU。Dice对目标和背景不平衡相对稳定IoU则更严格。在细线结构任务上我会额外记录HD95也就是预测边界到真实边界的95百分位距离用来判断边缘是否足够贴合而不仅仅是看重叠比例。4.2 对比实验结果模型相对参数量CT任务DiceCT任务IoU血管任务Dice血管任务IoU基础UNet100%0.8610.7980.8120.731UNetSE101.2%0.8730.8110.8230.744UNetCBAM103.8%0.8750.8140.8290.751UNetCA101.9%0.8800.8200.8390.763相对参数量是以基础UNet为100%归一化展示的不同实现会有细微差异但趋势是明确的。坐标注意力在两个任务上都比SE和CBAM高一点在细线结构上提升最明显Dice比基线高出2.7个百分点。这个幅度不算夸张但考虑到模块只增加不到2%的参数性价比已经很高了。还有一个值得关注的点SE和CBAM在CT大目标任务上差距不大。这也能理解当目标区域较大、边界信息可靠时通道注意力基本能满足需求。坐标注意力的主要受益场景就是细线、小病灶、边界模糊这些容易丢失位置信息的情况。4.3 从可视化中看到的变化除了指标我更关注注意力权重的可视化。把CA生成的a_h和a_w乘回特征图后能看到靠近目标中心以及沿目标轮廓方向的通道被明显放大。SE给出的权重往往是“整图均匀放大”而CA的放大区域更贴合目标的实际位置。在造影图像那种细长分支场景下宽度方向的权重能持续覆盖主干走向高度方向的权重能区分分支在纵向上的深度两者相乘后细线目标拿到了更高的响应。这种可视化同样能帮我判断模块到底有没有学到位置信息。如果训练后权重仍然很均匀那就意味着位置信号没有起作用需要检查reduction是否设置过大、训练迭代是否足够或者模块位置是否有问题。5. 实践中的调参与避坑指南5.1 位置不同结果天差地别我最早图省事只在最后一层编码器输出前加了一个CA结果提升很小有的任务甚至掉了0.5个点。后来把CA放到每个stage之后效果才稳定下来。原因是注意力模块如果只有一个相当于用同一个尺度的信息去做全局叠加无法在不同分辨率上分别进行位置校正。跳连接位置我也试过。对细线结构任务确实有用但对大器官分割任务反而会给浅层特征增加一些不必要的抖动。所以我的建议是先按“每个编码器stage 每个解码器stage”搭建。如果显存允许再单独尝试跳连接版本结合可视化决定是否保留。5.2 训练参数和经验性调整加了CA之后模型对学习率变得敏感一些。基础UNet在1e-3左右也能跑加了CA后我一般降到1e-4配合余弦退火调度效果更稳定。如果batch size不超过4BN的统计量容易抖动CA里的1×1卷积后面那个BN也容易受影响。这时候可以打开同步BN或者把CA内部的BN换成LayerNorm、GroupNorm实测后者在小batch下更稳。损失函数方面我建议用DiceLoss和BCE的混合再按任务加一点FocalLoss。医学分割常见的问题是前景占比很小如果只用BCE模型会倾向于把所有像素都预测为背景。CA能帮助聚焦目标区域但如果损失函数不平衡注意力也救不回来。5.3 常见问题速查表现象可能原因解决办法加了CA反而掉点模块位置不合理或reduction过大按编码器解码器stage放置reduction取16或32注意力权重可视化全均匀训练不充分或学习率过低增加迭代轮数调整学习率检查数据增强强度显存OOM在跳连接处加了过多CA分支只保留编码器和解码器CA或用梯度检查点细线目标没有明显提升下采样次数过多位置信息被抹掉在第一个stage后保留高分辨率特征或配合深监督CA参数不更新误接在冻结块后存在梯度截断检查网络结构确认CA参与正向梯度传播这些答案不一定对每个任务都成立但多数情况下按这个方向排查是高效的。坐标注意力不是万能药它最大的价值是给网络一个“方向感”。实际使用的时候要在自己的数据上多跑两组对比看看提升是不是来自偶然的随机种子。我个人在实际操作中的体会是加模块之前先想清楚目标场景最缺什么。如果你的分割任务本来就是大而清晰的目标加不加CA可能差别不大如果经常被小目标、细分支、边界模糊困扰CA是一个值得优先尝试的轻量方案。先跑通结构再看可视化再做消融比盲目堆一堆模块可靠得多。后面我准备继续把CA和深监督、多尺度特征融合放到一起试到时候再回来分享。