首页 / 资讯中心 / 文章详情

空间注意力模块实战指南:原理、代码与调参经验

空间注意力模块实战指南:原理、代码与调参经验 ★ FEATURED ARTICLE
做视觉任务这些年大家都会遇到一个很实际的需求网络深了、宽了但该学的东西没学会。SpatialAttention空间注意力模块是我调过的结构里性价比最高的一个——代码就几十行、参数量几乎可以忽略却能明显改变网络的关注重心。这篇文章不打算复读论文我想把在图像分类和目标检测里反复折腾这个模块的经验写出来它是怎么算的、放在哪一层最有效、有哪些坑我替你先踩过了。不管你是刚接触注意力机制的新手还是在为轻量模型提点发愁的工程师这个模块都值得花一下午研究明白。它适合作为你改造网络结构的第一个动手实验原理清晰、实现简单、效果可感而且就算调坏了对整体性能的影响也很有限拿它练手几乎没什么风险。1. 先搞清楚空间注意力在注意力机制里的“生态位”1.1 一张特征图里的“信息密度”差异先说个直觉。做分类任务的时候一张图里真正决定类别的区域往往只占一小块。比如一张猫的图片背景可能是沙发、地板、墙真正有用的就是猫那一坨。但卷积网络是滑动窗口式地处理整张图它对每个空间位置一视同仁计算资源被平摊掉了。这本身就是一种浪费更麻烦的是背景区域如果出现了和前景相似的纹理网络很容易被带偏。空间注意力要解决的就是这个问题给特征图的每个空间位置算一个权重重要的位置权重高、不重要的位置权重低然后用这个权重去重新缩放特征图。本质上是在做“位置选择”。这和注意力机制里的“Query-Key-Value”大框架不太一样它不搞复杂的相似度计算就是朴素的池化加卷积但效果出奇地好。这里有个容易被忽略的点空间注意力输出的不是一个标量也不是一个向量而是一张和输入特征图宽高一致的二维权重图。这意味着它对每个像素位置都有独立的加权决策。这个特性决定了它比全连接式的通道注意力更细腻也更适合中高层特征图上做精细化调整。1.2 为什么不直接用通道注意力很多初学者会把通道注意力和空间注意力搞混。通道注意力的代表就是SENet里的Squeeze-and-Excitation模块它对每个通道输出一个权重本质是“哪些语义模式重要”。而空间注意力输出的是二维空间权重本质是“哪些位置重要”。一个特别好的类比是通道注意力是拿着一张检查清单找特征比如“有没有毛茸茸的纹理”“有没有尖耳朵的形状”每个清单项打一个分空间注意力则是拿荧光笔在地图上圈重点区“这块区域值得多看”“那块区域可以扫一眼”。两者不是互斥的更多时候是互补的。CBAM这篇工作把两者串起来用先做通道注意力再做空间注意力效果比单独用任何一个都好。但要记住一点空间注意力并不依赖通道注意力存在它完全可以独立挂在任何主干网络上。很多任务场景下单独加空间注意力已经能带来不错的提升因为位置信息往往是模型最容易忽略的盲区。1.3 这个模块适合放在哪一层这是我在实际项目里花时间最多的问题。理论上看空间注意力可以插在网络的任何位置但不同位置的效果差别很大。浅层特征图分辨率高包含的是一些边缘、纹理、颜色信息。在这个阶段做空间注意力相当于让网络早期就去筛选像素级特征但实际上早期特征本身还没有形成语义概念筛选的意义不大反而增加了不少计算量。深层特征图分辨率低语义信息集中每个像素基本上已经代表某种语义区域了。这时候做空间注意力才真正在回答“这个语义区域重不重要”。我个人的经验是优先放在网络最后两个stage或者下采样之后。特征图分辨率掉到14×14或者7×7这种级别的时候空间注意力的效果最明显。如果放在第一个stage可视化注意力图会发现它高亮的几乎全是边缘轮廓对最终分类贡献很小计算开销却不小。2. 手写一个Spatial Attention模块从公式到代码2.1 最经典实现长什么样空间注意力有很多变体但最经典、最值得先搞懂的是CBAM里的实现。它的流程可以拆成三句输入特征图F沿通道维度分别做平均池化和最大池化得到两个二维特征图把这两个特征图在通道维度上拼接起来过一个卷积层把通道数压回1再接Sigmoid得到空间注意力权重。用公式表达就是M_s(F) σ( f^{7×7}( [AvgPool(F); MaxPool(F)] ) )其中f表示卷积核大小为7×7的卷积操作AvgPool和MaxPool都是沿着通道维做的得到的都是单通道的H×W特征图。从直觉上去理解它做的事情平均池化提取的是所有通道的“平均响应”代表了整体上哪些位置有信息最大池化提取的是所有通道的“最强响应”代表了哪些位置有最突出的特征。两边合在一起等于同时看了“大家怎么看”和“最敏锐的那个人怎么看”最后交给卷积去学一个最合适的组合方式。2.2 卷积核尺寸与感受野的取舍为什么默认用7×7的卷积核这是有讲究的。空间注意力要在二维平面上判断一个位置是否重要必须看到周围的上下文。如果卷积核太小比如1×1那它只能看单个像素根本不知道这个像素周围是什么判断“重要性”就成了无源之水。7×7的卷积核能覆盖更大范围让权重图更平滑、更有区域感。但大卷积核不是越大越好。我的实测经验是如果特征图分辨率本身已经不大了比如7×7一个7×7卷积的感受野已经覆盖了整张图和全局注意力没什么区别加上这个卷积的计算量占比会变得比较扎眼。所以在中低分辨率层我会把卷积核改成3×3或者5×5效果更稳计算也更省。给一个比较实用的选型建议插在stage 3之前用5×5或者7×7插在stage 4之后用3×3。如果你用的是轻量网络比如MobileNet系列建议直接默认3×3因为大卷积核在深度可分离卷积为主的架构里会破坏计算量的平衡。2.3 平均池化与最大池化缺一不可可能有人会问既然已经有了平均池化为什么还要加最大池化这其实是我一开始也不理解的地方。纯用平均池化的话空间注意力会更关注“广泛的、持续的响应区域”权重图通常比较柔和但容易把人眼觉得重要的强特征给平滑掉。纯用最大池化的话权重图会更锐利强响应区域被高亮但训练初期噪声大的时候最大池化很容易被异常点带偏。把两个池化结果拼起来再卷积等于让网络自己决定该信哪一边这比人工定一个规则要灵活得多。我自己的解释方式是把两种池化比作两个“顾问”平均池化是个稳健派告诉你哪里整体上都有响应最大池化是个激进派告诉你哪里响应最强、最值得看。两个顾问的意见都有道理最后听谁的卷积层会学会。2.4 完整PyTorch实现直接贴一份可以跑的实现这段代码我在PyTorch 1.8以上版本都验证过import torch import torch.nn as nn class SpatialAttention(nn.Module): def __init__(self, kernel_size7): super(SpatialAttention, self).__init__() assert kernel_size in (3, 7), kernel_size must be 3 or 7 padding kernel_size // 2 self.conv nn.Conv2d(2, 1, kernel_size, paddingpadding, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out torch.mean(x, dim1, keepdimTrue) max_out, _ torch.max(x, dim1, keepdimTrue) attention torch.cat([avg_out, max_out], dim1) attention self.sigmoid(self.conv(attention)) return x * attention几个细节说一下卷积层没有用bias因为后面接Sigmoid偏置对权重的零点位置影响不大省掉能少一点参数。输入形状是[B, C, H, W]输出形状和输入完全一致方便直接接到任何位置。torch.max(x, dim1, keepdimTrue)返回两个值我们只需要第一个也就是max_out。有些实现会在卷积之后加一个BatchNorm我试过效果没有显著提升反而在批大小较小的时候引入额外抖动所以默认不加。3. 接入主干网络时最常踩的三个坑3.1 模块插到哪里才不破坏残差结构接法不对再好的模块也是白搭。最常犯的错误是把空间注意力插在残差相加之后、ReLU之前还是之后其实都可以但要理解不同位置的语义。以ResNet为例比较稳妥的接法是放在Block内部输出之后、残差相加之前。这样空间注意力先对Block学到的特征做了一次位置加权再和恒等映射相加。如果注意力把某些位置压到接近0残差分支的这一路信息被抑制了但shortcut那一路仍然保留原始信息梯度仍然能顺畅回传。如果放在残差相加之后就相当于对“已经融合了shortcut信息的输出”再做加权。这时候注意力要同时处理“新学到的信息”和“保留下来的老信息”任务更重而且如果权重接近0整个Block的输出会被压制梯度虽然还能走shortcut但主分支的信息就废掉了。我在实际项目中比较倾向于第一种SpatialAttention加在残差分支末端也就是常规的conv - bn - relu - conv - bn后面再执行相加。这也是CBAM论文里推荐的挂载位置绝大多数主干网络都不需要改结构就能插进去。3.2 通道注意力与空间注意力顺序千万别搞反如果要把通道注意力和空间注意力串起来用顺序是有讲究的。CBAM作者做过消融实验结论是通道注意力在前、空间注意力在后效果最好。这个顺序不是玄学背后有逻辑通道注意力先做相当于先筛选“哪些语义模式值得关注”把不重要的通道压下去空间注意力再做相当于在已经筛选过的特征图上做“位置选择”。这样空间注意力只用关心位置而不用分心去对抗那些无关通道带来的噪声。如果反过来空间注意力得在包含大量无关通道信息的特征图上算位置权重效果很容易被噪声带偏。我自己的经验还发现如果空间注意力在前训练初期loss曲线的波动会明显更大收敛也慢一些。这可能是因为空间权重先对高维特征做了乘法缩放再走通道注意力时通道维的统计量被扰乱了。3.3 训练策略与超参调整加了空间注意力之后训练设置不一定非得大改但有几个细节值得注意。第一学习率。空间注意力的参数非常少一个7×7卷积只有66个参数对整体收敛速度的影响很有限。所以一般情况下不需要单独给这个模块设学习率。但如果你发现loss下降变慢可以先怀疑是不是模块放多了而不是调学习率。第二初始化。默认的Kaiming初始化就够用。我试过给最后一层卷积做特殊初始化比如初始化为0让注意力一开始接近全1相当于先让网络按原来的方式训练再慢慢学注意力。这个思路是对的但默认初始化其实已经很接近这个状态了不需要额外操作。第三正则。空间注意力天然带有一种特征选择性质它等效于给特征图乘了一个随输入变化的mask这在某些任务上会带来轻微的正则效果。如果你原来用了比较重dropout加了模块之后可以试着把dropout比例降一点效果可能反而更好。第四可学习缩放。这是我个人比较喜欢的一个小技巧在Sigmoid之后乘一个可学习的标量初始值设为1。class SpatialAttentionV2(nn.Module): def __init__(self, kernel_size7): super().__init__() padding kernel_size // 2 self.conv nn.Conv2d(2, 1, kernel_size, paddingpadding, biasFalse) self.sigmoid nn.Sigmoid() self.scale nn.Parameter(torch.ones(1)) def forward(self, x): avg_out torch.mean(x, dim1, keepdimTrue) max_out, _ torch.max(x, dim1, keepdimTrue) attention torch.cat([avg_out, max_out], dim1) attention self.sigmoid(self.conv(attention)) * self.scale return x * attention这个小改动让模块在训练初期可以自动决定注意力的“强度”我实测在多个数据集上能让收敛更稳定一点点。4. 加了模块没效果这份排查清单值得先看4.1 加了模块反而掉点大概率是这几种原因我最开始用Spatial Attention的时候也遇到过加了反而掉点的情况。后来排查多了发现原因基本集中在以下几类。任务太简单、模型已经过拟合。如果模型本身很小、任务很轻松加注意力等于额外引入了一组非线性变换虽然参数量小但它干扰了原有特征分布。这种情况下的表现是训练集上loss下降更快验证集上精度反而低。解决办法是配合更强的数据增强或者更大的dropout。模块放的位置太浅。前面提到过浅层特征还没形成语义空间注意力在浅层更像是在做边缘提取没什么帮助甚至有害。如果你的网络比较深只在最后一个stage加别每层都加。加的层数太多。每个Block后面都加空间注意力会导致特征被反复“高亮”网络变得过于聚焦反而丢了全局信息。比较好的做法是隔几个Block加一次或者只在最后两个stage加。这个是我反复验证过的经验。Sigmoid饱和。如果注意力图的数值长时间接近0或接近1说明卷积层输出的绝对值过大Sigmoid进入了饱和区。梯度传不回去模块等于白加。解决方法是给卷积层做权重初始化时把增益调小一点或者加一个LayerNorm稳定一下数值范围。4.2 显存和计算量到底焦虑不焦虑很多人一听“注意力”就害怕显存爆炸这里得分清楚是哪一种注意力。经典的Spatial Attention模块和Self-Attention那种计算量完全是两码事。经典Spatial Attention增加的计算量主要来自那个7×7卷积输入是2通道、输出是1通道卷积核尺寸7×7对于一张H×W的特征图计算量大概是2×7×7×H×W次乘加。对比主干网络动辄上亿次的FLOPs这个增加量基本可以忽略。显存上也只是多存了一张注意力权重图和两个池化结果分辨率还不高完全不用慌。真正吃显存的是Non-local风格的空间注意力它对特征图里每个像素都要和所有像素算相关度复杂度是O(N^2)的NH×W。如果特征图是28×28N784注意力图就有将近61万个元素继续往下不用算显存早就没了。我把这两类做了一张对比表类型通道特征核心操作计算复杂度适用场景经典Spatial Attention池化压缩卷积O(H×W)通用主干网络开销可忽略Non-local空间注意力矩阵乘法相似度计算O(H×W)²小分辨率、长距离依赖场景做工程部署或者模型压缩的时候首选一定是经典Spatial Attention它便宜、有效、不容易出问题。4.3 怎么判断注意力图学了没学判断模块有没有学到东西最直接的办法是可视化注意力图。把Sigmoid输出的权重图resize到和原图一样大然后用热力图叠加到输入图像上立刻就能看出端倪。正常的注意力图应该呈现出明显的区域聚焦比如分类任务里高亮的是目标主体而不是背景检测任务里高亮的是目标附近而不是整张图均匀变亮。如果注意力图整片都是接近1的灰色说明Sigmoid输出饱和了模块退化为一个近似恒等变换如果注意力图全是接近0的黑色说明方向梯度出了偏差需要检查卷积层初始化。如果注意力图高亮区域在训练过程中频繁跳动从前景跳到背景又跳回来多半是训练不稳定可以考虑把学习率降低或者把Sigmoid换成tanh配合缩放。我在一个检测任务里遇到过注意力图开始高亮目标边缘后来慢慢扩展到了目标内部这是一个正常的学习过程耐心等几个epoch就好。可视化代码很简单只要把注意力权重图保存下来用matplotlib里的imshow配合cmapjet就能出图。关键是不要等到整个训练结束再看训练早期每几百个iteration就抽查一次能更快发现问题。5. 从Spatial Attention延伸出去的高级玩法5.1 Non-local式空间注意力为什么又慢又“稳”前面提到了Non-local注意力它算是空间注意力在长距离依赖方向上的极致延伸。核心思路是对特征图上的每个位置计算它和所有其他位置的相似度然后用这个相似度对所有位置的特征做加权求和相当于每个像素都能“看到”全图。这个机制的优点很明显能捕捉远程依赖特别适合大目标、全局结构明显的任务。但它的问题也很致命——计算量是O(N^2)的N是空间位置数。哪怕输入特征图只有28×28也要算784×784的相似度矩阵显存直接告急。对这种模块我的使用经验是要严格控制输入分辨率。一般只在高层的、已经下采样很多倍的特征上使用低于14×14才比较可行。如果要在更高分辨率上用就得配合分块策略或者近似注意力但这已经是另外一个话题了。5.2 轻量化改造把全局建模压缩到线性复杂度工程上更实用的方案是把空间注意力改造成线性复杂度。一个没有太多理论负担、但实测有效的做法是把全局信息先压缩成一个全局向量再和局部特征交互。具体思路是对空间维度做全局平均池化得到每个通道的全局描述再把全局描述通过一个小网络变换后加到原始特征图的空间维度上。这样每个位置都间接看到了全局信息但计算复杂度只有O(H×W)高分辨率下也跑得动。GCNet相关的工作基本就是这个路子。如果你只是想在移动端模型里塞一个空间注意力我建议先用经典Spatial Attention不要直接上这种全局交互方案。经典方案在部署时已经足够轻了再复杂收益就变小了。5.3 注意力可视化带来的调试思路最后说一个我非常推荐的调试习惯把空间注意力图融入你的日常跑图流程里。每次训练到一定阶段挑几张训练样本把注意力权重和原图叠起来看。这不是为了发论文凑图而是为了快速定位问题。比如分类模型效果不好如果注意力图高亮区域和真实目标对不上说明模型提取的特征本身有问题问题可能出在数据或主干网络而不是注意力模块本身。如果注意力图高亮区域很准但最终精度还是上不去那瓶颈可能在分类头就需要去调输出层的结构。这个习惯帮我在很多任务里少走了弯路因为注意力图往往比loss曲线更直观。loss只是一条曲线注意力可视化则能直接把模型的“审美偏好”摆在你面前。我自己用得最顺手的组合是主干网络最后两个stage挂经典Spatial Attention核大小用5×5Sigmoid前加一个可学习缩放因子再接一个轻量通道注意力。这个配置在多个图像分类和目标检测任务里都表现得很稳定而且代码改动量非常小。如果你正准备改造自己的网络结构不妨从这个模块入手。它不是最花哨的却是最容易亲身体会到“注意力机制到底在做什么”的一个。
阅读完成 · 觉得有帮助?
咨询建站