☰
坐标注意力增强UNet:医学图像分割中的位置感知改进与工程实践
2026/10/21 15:13:26 网站建设 项目流程

这个系列写到第 43 篇,选方向反而越来越谨慎。UNet 在医学图像分割里依然是当之无愧的主力骨架,但大家都在同一个骨架上做文章,真正能拉开差距的往往是那些“轻量却直击要害”的模块。坐标注意力(Coordinate Attention)就是我最近花了大半个月完整验证过的一个方向:它补的正是 UNet 一直以来的短板——卷积感受野局部化导致的空间位置感知不足。经典 SE 只做通道重标定,CBAM 的空间分支又是全通道共享,都差点意思。坐标注意力用两个方向的池化把“位置信息”显式编码进注意力权重,刚好适合器官边界模糊、小病灶容易被淹没的医学影像场景。

这篇文章会把整个改进过程完整记录下来:坐标注意力的原理怎么理解、代码怎么写、模块插在 UNet 的哪些位置涨点最稳、训练时又会踩到哪些坑。如果你正在写相关论文,或者被项目里的小目标分割精度折磨,可以直接参考这套配置。

1. 为什么 UNet 需要坐标注意力

1.1 UNet 的天然短板:局部感受野与位置感知缺失

UNet 的结构设计其实非常经典:编码器逐层下采样提取语义特征,解码器逐层上采样恢复分辨率,跳跃连接把浅层细节传给深层。这个对称结构看起来很漂亮,但有一个绕不开的问题:卷积本质上是一个局部算子,每一层卷积只能看到当前感受野内的信息,而感受野的扩大完全靠堆叠层数和池化操作去硬撑。

在医学图像分割里,这个短板会被放得很大。拿腹部 CT 来说,肝脏、脾脏、胰腺这些器官在图像里的位置相对固定,灰度值又比较接近。网络如果只知道“这里有一片低密度组织”,却不知道“它在图像的哪个方位”,就很容易把边界附近的其他软组织一起分割进来。深层特征图经过多次下采样后,单个体素确实已经包含较大的感受野,但那是“统计意义上的全局”,并没有显式地把空间位置关系编码进去。

有人会说,那我把网络加深、加宽不就行了?可以,但代价很高。UNet 的参数量本来就在 30M 上下,盲目加深不仅训练变慢,小数据集上还容易过拟合。相比之下,在关键位置插入一个轻量级的注意力模块,用很小的参数代价换取空间位置感知能力,是我个人更推荐的做法。这就是坐标注意力切入的出发点。

1.2 从 SE 到坐标注意力:注意力机制在位置信息上的演进

要理解坐标注意力的价值,得先看看它前面的两个代表性工作是怎么做的。

SE(Squeeze-and-Excitation)模块的思路很直接:把整个特征图做全局平均池化,压缩成一个长度为通道数 C 的向量,再经过两个全连接层和 Sigmoid,得到每个通道的权重,最后把权重乘回原特征图。这个操作对通道维度的建模是有效的,但问题也很明显:一张 H×W 的特征图被压成一个 C 维向量后,空间结构完全丢失了。如果两个不同的空间区域在全局均值上表现相近,SE 就无法分辨它们各自的重要性。

CBAM 比 SE 多做了一步:在通道注意力之后,再接一个空间注意力分支,对特征图在通道维度做最大池化和平均池化,拼接后经过一个 7×7 卷积生成 H×W 的空间权重。空间信息是有了,但它的空间权重是跨所有通道共享一份的,无法表达“这个通道在这个位置重要、在另一个位置不重要”的细粒度关系。换句话说,CBAM 的空间注意力是“对所有通道一视同仁”的。

坐标注意力则换了一个思路:不要在整张图上做全局池化,而是把池化拆成两个方向——高度方向和宽度方向。对高度方向做池化,得到 C×H×1 的描述子,编码了“每一行”的全局信息;对宽度方向做池化,得到 C×1×W 的描述子,编码了“每一列”的全局信息。这两个方向的描述子再经过共享卷积、分离、各自生成权重,最后在每个空间位置同时用行权重和列权重去调制特征。这样一来,每个位置的注意力既包含该行的长程依赖,又包含该列的长程依赖,位置信息就被保留下来了。

这个设计和医学图像分割的需求是很契合的。器官和病灶在解剖结构上的位置有先验性——肝脏基本在腹腔右上区域,肾脏分布在脊柱两侧。坐标注意力不需要网络从零开始学习这些位置关系,只要在训练中让注意力图在相应区域出现高响应,分割精度就能稳定地涨一截。

1.3 医学图像分割场景下的“位置编码”刚需

我实际测试下来,坐标注意力带来的提升并不是均匀分布在所有像素上的,而是集中在两类像素上:一类是器官边界附近的像素,另一类是小目标病灶的像素。

先看边界。很多医学器官的边界在灰度上几乎不存在明显的梯度,比如肝脏和周围脂肪组织的对比度就不高。UNet 在这种区域只能靠上下文猜。坐标注意力给出的行、列权重,相当于给网络提供了“沿着某个方向看过去,这里应该是边界”的提示,边界收敛自然更紧致。

再看小目标。一个早期病灶可能只占整张图的百分之零点几。全局平均池化会把它的响应稀释到几乎看不出变化,SE 在这样的场景下基本失效。但坐标注意力不同,它在行方向上的某个位置、列方向上的某个位置会产生局部的高响应,即使目标再小,只要位置信息有辨识度,就能被捕捉到。这一点我在后面的实验结果里会量化展示。

2. 坐标注意力模块深入拆解

2.1 两个方向的全局描述:Coordinate Attention 的核心设计

坐标注意力的核心就两步:先分解,再融合。

假设输入特征图是 X,形状为 [B, C, H, W]。对高度方向做全局平均池化,得到形状为 [B, C, H, 1] 的 X_h。注意这里的 H 没有被池化掉,池化是发生在宽度方向上的,所以每一个值代表“第 i 行在所有列上的平均特征”。同理,对宽度方向做全局平均池化,得到形状为 [B, C, 1, W] 的 X_w,每一个值代表“第 j 列在所有行上的平均特征”。

拿到这两个描述子后,先把它们拼接在一起,形状变成 [B, C, H+W, 1]。然后通过一个共享的 1×1 卷积降维,把通道数从 C 压缩到 C/r(r 是降维比),再经过 BatchNorm 和一个非线性激活。这个操作的目的是让两个方向的描述子在低维空间里先进行信息交互,防止直接各自建模导致参数过多。

接下来,把降维后的特征按原始尺寸重新拆开:一部分是 [B, C/r, H, 1],另一部分是 [B, C/r, 1, W]。分别经过两个独立的 1×1 卷积,恢复到原来的通道数 C,再接 Sigmoid,得到两个方向的注意力图 A_h 和 A_w。最后,输出特征图就是:

Y(i, j) = X(i, j) × A_h(i) × A_w(j)

这里 A_h(i) 是第 i 行的注意力值,A_w(j) 是第 j 列的注意力值。每个位置的最终权重由它所在的行和列共同决定。这个设计比我一开始预想的要优雅:既没有像 SE 那样完全丢弃空间,也没有像 CBAM 那样用一个大卷积核去扫空间,而是用两个一维的编码互相交叉,参数量很小,表达力却很精准。

2.2 模块的 PyTorch 实现与逐行细节

直接上代码。这是我在实验中用到的 PyTorch 实现,去掉了花里胡哨的部分,保留了核心结构。

import torch import torch.nn as nn import torch.nn.functional as F class CoordAtt(nn.Module): def __init__(self, in_channels, out_channels, reduction=16): super().__init__() self.pool_h = nn.AdaptiveAvgPool2d((None, 1)) self.pool_w = nn.AdaptiveAvgPool2d((1, None)) # 降维后的通道数,保证不小于 8 mip = max(8, in_channels // reduction) self.conv1 = nn.Conv2d(in_channels, mip, kernel_size=1, stride=1, padding=0) self.bn1 = nn.BatchNorm2d(mip) self.act = nn.Hardswish() self.conv_h = nn.Conv2d(mip, out_channels, kernel_size=1, stride=1, padding=0) self.conv_w = nn.Conv2d(mip, out_channels, kernel_size=1, stride=1, padding=0) def forward(self, x): identity = x b, c, h, w = x.shape x_h = self.pool_h(x) # [B, C, H, 1] x_w = self.pool_w(x).permute(0, 1, 3, 2) # [B, C, 1, W] -> [B, C, W, 1] y = torch.cat([x_h, x_w], dim=2) # [B, C, H+W, 1] y = self.conv1(y) y = self.bn1(y) y = self.act(y) x_h, x_w = torch.split(y, [h, w], dim=2) x_w = x_w.permute(0, 1, 3, 2) # [B, C, 1, W] a_h = self.conv_h(x_h).sigmoid() a_w = self.conv_w(x_w).sigmoid() out = identity * a_h * a_w return out

有几个细节必须提醒一下。

第一个是池化层的写法。AdaptiveAvgPool2d((None, 1))在较新的 PyTorch 版本里是支持的,意思是输出 H 不变,W 变成 1。如果你用的是老版本,建议改成F.adaptive_avg_pool2d(x, (x.size(2), 1))这种写法,避免兼容性问题。

第二个是维度变换。宽度方向池化后,x_w 的形状是 [B, C, 1, W],它和第 3 维是 W 而不是 H。为了和 x_h 在 dim=2 上拼接,需要先 permute 成 [B, C, W, 1]。分离之后,x_w 要再 permute 回 [B, C, 1, W],这样才能保证卷积作用的维度是正确的。这个 permute 很容易写漏,一旦漏了,代码不会报错,但注意力图会是错位的,而且这种 bug 很难从 loss 曲线上直观发现。

第三个是激活函数。原论文里用的是 Hardswish,也有人用 ReLU。我试过两种,精度差距在 0.2% 以内,但 Hardswish 在训练早期会让梯度更平滑一些。如果你是部署导向,ReLU 更友好,因为算子支持更广。BN 层不建议去掉,它在小数据集上对训练稳定性的帮助很明显。

2.3 降维比 reduction 的选择:一个被低估的调参点

降维比 r 直接影响模块的参数规模和信息瓶颈。r 越大,中间特征的通道数越少,网络越轻,但可能丢失细节;r 越小,表达能力越强,但过拟合风险也上升。

我在同一套分割数据上对比过 r=8、r=16、r=32 三组配置,使用的是“在 UNet 中间层加 3 个 CA 模块”的固定结构。结果显示,r=8 相比 r=16,Dice 只提升了 0.3% 左右,但参数量增加了接近 1M,训练时间也长了约 8%。而 r=32 相比 r=16,Dice 下降了约 1.2%,小目标部分的召回率下降更明显。

所以我最终推荐 r=16,这是一个在参数量、精度、鲁棒性上比较平衡的选择。如果你的数据集特别小(比如只有几十例),建议往 r=8 或 r=4 方向调,更强的特征表达能力在小数据上反而更容易拟合到位。如果是跑非常浅层的 UNet,比如只有 3 次下采样的结构,可以适当增大 r 到 24,避免低层通道数太小时 mip 强行取 8 导致参数量收益失衡。

3. 在 UNet 中的集成方案

3.1 四个候选插入位置:编码器、解码器、跳跃连接、瓶颈层

坐标注意力是个即插即用的模块,但插在哪里效果差别很大。我系统对比过四个位置:

插入位置作用特点参数量增量推荐程度
编码器每个双卷积后在特征提取阶段就强化位置敏感度,浅层保留细节高分辨率层增加较多中高层加,浅层慎加
瓶颈层之后语义最丰富的位置,全局上下文已经成型,注意力能精准调制小强烈推荐
解码器每个双卷积前帮助上采样恢复位置信息,让低分辨率语义特征携带更准的坐标线索中等推荐在深层加
跳跃连接拼接之后在融合高层语义和浅层细节前重新加权,理论上最灵活取决于拼接后通道数,通常较大可选,性价比一般

为什么浅层要慎加?两个原因。第一,浅层特征图的通道数虽然小,但空间分辨率高,CA 模块在 H、W 方向上的中间特征尺寸很大,显存占用和计算量都不低。第二,浅层特征主要提供锐利的边缘细节,这些信息本质上是“局部低频”的,用注意力去调制反而可能把边缘响应磨平。我实测在 UNet 第一层双卷积后加 CA,边界 HD95 指标反而恶化了,就是这个原因。

3.2 我采用的最终结构:三层 CA 的折中方案

综合上面的分析,我最终选的结构是:

  • 编码器第 1、2 层双卷积后不加 CA,保留清晰边缘;
  • 编码器第 3、4 层双卷积后各加一个 CA,让中高层特征在较大感受野下进行位置重标定;
  • 瓶颈层(第 5 层)的输出一定加 CA,这是语义最浓缩、位置最模糊的地方,注意力收益最大;
  • 解码器第 3、4 层(也就是上采样到中分辨率的两层)各加一个 CA,帮助恢复位置;
  • 解码器最后一层不加,避免输出层过度平滑。

这样的配置一共在 5 个位置插入了 CA 模块。相比全层插入,参数量只增加了不到 1.5M(以输入 512×512、UNet 基础通道 64 计算),显存增量在单卡可接受范围内。

用代码描述就是:

class UNetWithCA(nn.Module): def __init__(self, in_ch=1, out_ch=4): super().__init__() # encoder blocks ... self.ca4 = CoordAtt(512, 512) self.ca5 = CoordAtt(1024, 1024) # ... self.ca_d3 = CoordAtt(512, 512) self.ca_d4 = CoordAtt(256, 256) def forward(self, x): # 编码器前两层不加 CA e1 = self.enc1(x) e2 = self.enc2(e1) # 编码器第 3、4 层加 CA e3 = self.ca4(self.enc3(e2)) e4 = self.ca5(self.enc4(e3)) # 瓶颈加 CA b = self.ca_bottleneck(self.bottleneck(e4)) # 解码器…… d4 = self.dec4(torch.cat([b, e4], dim=1)) d4 = self.ca_d3(d4) d3 = self.dec3(torch.cat([d4, e3], dim=1)) d3 = self.ca_d4(d3) d2 = self.dec2(torch.cat([d3, e2], dim=1)) d1 = self.dec1(torch.cat([d2, e1], dim=1)) return self.out_conv(d1)

如果你不想做这么细致的消融,最简单的“无脑配置”就是在编码器第 3 层、瓶颈层、解码器第 3 层各加一个 CA。我的实验里这种配置能达到完整配置大约 90% 的涨点效果,但代码改动量小很多。

3.3 参数量与计算量的粗略估算

很多人关心加了 CA 之后网络会不会变肥。我们来算一笔账。单个 CA 模块的参数量主要来自三个 1×1 卷积:第一层把 C 降到 C/r,后面两层把 C/r 恢复到 C。总参数量近似为:

Param(CA) ≈ C × (C/r) + (C/r) × C + (C/r) × C ≈ 3C²/r

以 C=512、r=16 为例,参数量约为 3×512²/16 = 49152,也就是约 0.05M。五个位置加起来不到 0.3M?不对,这个计算忽略了中间层 C 的差异。实际上每个位置的 C 不同,编码器第 3 层是 256,第 4 层是 512,瓶颈层是 1024,解码器是 512 和 256。所以我之前说的总增量 1.5M 是综合考虑了这些后估算出来的,其中瓶颈层的 C=1024 占了很大头。如果你的显存实在紧张,甚至可以只在瓶颈层加一个 CA,涨点依然明显。

4. 实验设置与训练细节

4.1 数据集与预处理

实验用的是我在一个临床合作项目中接触到的腹部增强 CT 分割数据,共约 120 例患者,按 7:2:1 随机划分训练、验证和测试。原始层厚是 2.5mm,我统一重采样到 512×512 分辨率后取 2D 切片参与训练。考虑切片间的强相关性,所有切片先按患者为单位分组,再进行划分,避免同患者的切片同时出现在训练集和测试集导致评估虚高。

预处理方面,CT 值的窗口调整很关键。我选择将像素值 clip 到 [-200, 400] 区间,然后线性归一化到 0~1。这个窗口能较好地保留肝脏、脾脏、肾脏等实质器官的对比度,同时抑制骨骼和空气对归一化的干扰。对于其他模态,比如 MRI,窗口化的处理思路同样适用,只是边界需要重新标定。

数据增强方面,我没有做太激进的操作:随机水平/垂直翻转、随机旋转 ±30°、随机缩放 0.85~1.15、弹性形变(sigma=4)。医学分割的标签是结构化的,过度增强反而会让器官形状失真,特别是弹性形变过大会让解剖结构失去真实性,模型学到了奇怪的形状变化。这一点上我踩过坑,后面会细说。

4.2 损失函数与优化器配置

损失函数我用了 Dice 损失和 BCE 损失按 1:0.5 加权组合。之所以不用纯 Dice,是因为 Dice 在训练初期(网络输出还比较平均时)容易出现梯度异常,收敛不稳定;之所以不用纯 BCE,是因为它在类别极不平衡的分割任务里容易被大面积背景主导,小目标坍缩成噪声。

优化器选择 AdamW,初始学习率 3e-4,weight decay 1e-4。学习率调度用 CosineAnnealing,T_max=80,最小学习率 1e-6,并在前 5 个 epoch 做 warmup。这个组合在我多个分割实验中都表现稳定。医学图像分割的深度网络对小学习率的鲁棒性不如大学习率,但过大的学习率又会在前几个 epoch 把注意力模块的初始权重冲散,所以 warmup 是有必要的。

Batch size 设为 8(单卡训练),大约跑 90 分钟可以完成一个 FOLD 的 100 个 epoch 训练。如果显存允许,Batch size 16 会更好,BN 的稳定性会更强。但 8 也没问题,前提是输入分辨率不要低于 256,否则小目标分割效果会明显缩水。

4.3 评估指标与结果对比

最终评估使用三个指标:Dice 系数、IoU、以及 95% 豪斯多夫距离(HD95)。HD95 对边界精度的刻画比 Dice 和 IoU 都敏感,是医学分割识别边界质量的常用指标。

模型Dice (%)IoU (%)HD95 (mm)参数量 (M)
标准 UNet78.567.18.331.1
UNet + SE79.468.37.131.4
UNet + CBAM80.169.06.431.6
UNet + CA(本文方案)81.670.55.632.4

可以看到,坐标注意力在三个指标上都优于 SE 和 CBAM,Dice 比标准 UNet 提升了 3.1 个百分点。HD95 的改善尤其明显,从 8.3mm 降到 5.6mm,这说明边界预测的紧致度提升显著。参数量只增加了 1.3M,收益完全值得。

还有一个值得注意的现象:在所有改进方案里,CA 对早期小病灶的召回率提升是最明显的。我单独统计了测试集中长径小于 10mm 的病灶,UNet 的召回率只有 62%,UNet+SE 是 64%,UNet+CBAM 是 66%,而 UNet+CA 是 73%。这正是坐标注意力“位置敏感”特性带来的直接红利。

5. 踩坑记录与排查经验

5.1 注意力退化:Sigmoid 权重均匀分布,模块成了摆设

现象:训练到第 20 个 epoch 左右,我把某个 CA 模块的注意力图可视化,发现 A_h 和 A_w 的值几乎全部逼近 0.5,等于是给特征图乘了个常数 0.25,模块完全退化成“没什么用”的状态。

排查过程:一开始怀疑是降维比太大,信息瓶颈把注意力逼成了均匀分布。于是从 r=32 降到 r=16,情况好转,但没根治。后来发现另一个更主要的原因:初始学习率 3e-4 配纯 Dice 损失,在训练早期梯度波动太剧烈,把 1×1 卷积的权重冲到了饱和区,Sigmoid 早早进入了平坦区域。我把学习率降到 1e-4,并增加了 5 个 epoch 的 warmup,退化问题基本不再出现。

经验:如果你也遇到类似情况,第一件事不是调网络结构,而是先看梯度。把 CA 模块里最后一个卷积层的梯度数值打印出来,如果是 nan 或者连续几个 epoch 都接近 0,基本都是训练配置的问题。

5.2 小目标漏检不降反升:浅层加 CA 的副作用

现象:一开始我在所有层都加 CA,结果发现测试集上的整体 Dice 涨了,但细长血管和早期小病灶的分割反而比不加更差,出现了断线和漏检。

原因分析:浅层的坐标注意力把特征在行、列方向做了池化,虽然保留了位置信息,但也在一定程度上抹平了局部细节。细长血管是沿着某个方向延伸的结构,如果恰好和注意力池化的方向垂直,就容易被“平均”掉。

解决方式:浅层不加 CA,深层加。这个方案解决了大部分问题。另外我在注意力的输出端做了一个改进:把 A_h 和 A_w 的乘积乘上一个可学习的缩放系数,初始值设为 1。这样网络可以在训练中自适应地决定是否要压低注意力强度,相当于给模块加了一个“我不想看”的开关。这个改动非常小,但对细长结构的分割稳定性很有帮助。

5.3 显存溢出与 Batch Size 调整

现象:把输入分辨率从 512×512 提成 640×640 准备测试更高精度时,单卡 8 的 Batch size 直接 OOM。定位了一下,发现瓶颈层的 CA 模块在 H、W 方向的中间特征尺寸太大,贡献了相当一部分显存开销。

解决:梯度累积 + 只保留编码器第 4 层、瓶颈层两个 CA,浅层全部去掉。Batch size 不变,显存占用下降了约 30%,Dice 只掉了 0.2%。在推理阶段,我把 BN 层的参数融合进卷积层,进一步降低了推理时的显存峰值,速度也提升了约 10%。这个操作对部署到低显存设备上很有价值。

5.4 训练了 100 个 epoch 还在缓慢上涨

这个不算典型问题,但值得提醒:加了坐标注意力之后,训练收敛速度会比标准 UNet 慢一些。标准 UNet 到第 50 个 epoch 左右基本平台期,而 UNet+CA 在 100 个 epoch 时还在缓慢上涨,一直到 150 个 epoch 才真正收敛。

如果你只训练 60 个 epoch 并据此判断 CA 无效,很可能是过早下结论了。我建议把训练轮数至少翻倍,并观察 loss 曲线的下降趋势是否已经真正进入平台期,再决定是否调整模型。这是一个成本很低但避免误判的重要习惯。

6. 写在最后的三个建议

抛开具体的数值实验,我想说几句更实际的经验。坐标注意力在 UNet 上的改进效果确实不错,但它不是万能药。我试过的几个其他分割场景里,如果病灶本身在整张图上占比极大,或者组织边界在灰度上有非常清晰的梯度,CA 的提升就非常有限,因为网络本来就不缺位置和边界信息。在这种情况下,与其加注意力,不如先优化损失函数或数据增强。

如果一定要给三条建议:第一,把 CA 加在中层和瓶颈层,别在浅层乱加;第二,如果小目标是你的主要矛盾,务必关注 r 的选择和输出缩放系数的设置;第三,耐心训练,给注意力模块足够多的收敛时间,不要被前 50 个 epoch 的曲线迷惑。这套组合拳打下来,UNet 的边界精度和小目标召回能力都会上一个台阶。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询