
最近在读语义分割方向的论文正好看到Dual Attention这篇读完之后有种“原来还可以这么直接算”的感觉。所谓Dual Attention核心是把注意力机制拆成两个维度分别处理一个管空间位置之间的关联一个管通道之间的关联最后把两路结果融合回特征图。这个思路在场景理解、分割、检测里都很有代表性尤其是那些需要远距离上下文信息的任务。这篇笔记我按自己的阅读习惯整理出来包含论文思路拆解、模块的原理解释、PyTorch复现代码还有踩过的一些坑。如果你在做计算机视觉方向尤其是分割、场景理解、注意力机制相关的工作这篇内容应该对你有用。1. 这篇论文到底想解决什么问题1.1 CNN的软肋局部感受野的局限性常规卷积神经网络处理图像时每个卷积核看到的范围非常有限。哪怕堆叠很多层卷积感受野也是线性增长而且中间层提取的多是局部纹理和局部结构很难直接建立两个相距很远的像素之间的关联。对语义分割来说这就成了很现实的问题一辆被树干遮挡的车车头和车尾在图像里可能隔了几百个像素但它们是同一个物体一个骑行者身体和旁边的行人颜色相近但语义完全不同。网络如果只盯着局部邻居很容易把同一物体的不同部位分成两半或者把不同物体粘连在一起。早期解决办法是空洞卷积它能在不增加参数的情况下扩大感受野但空洞卷积的采样点是稀疏的依然存在“网格效应”——它只关心特定间隔的采样位置对全局关系的建模本质上还是间接的。另一个常见思路是ASPP这类多分支池化结构用不同膨胀率的卷积去抓多尺度信息但它面对的是“尺度”问题不是“长距离依赖”问题。碰巧长距离依赖、上下文一致性这些痛点正好是注意力机制擅长的事情。1.2 两条互补的注意力分支整体设计思路Dual Attention的整体结构说起来很清晰从backbone拿到一个特征图之后把它同时送入两个并行的注意力分支。第一个分支叫位置注意力模块负责计算每个空间位置对其他所有位置的响应权重生成一个尺寸为N×N的注意力矩阵其中N等于特征图的高乘宽。第二个分支叫通道注意力模块负责计算每个通道对其他所有通道的响应权重生成C×C的矩阵其中C是通道数。两个分支的输出和原始特征做残差相加再合并在一起。这种设计最巧妙的地方是把特征图的两个维度拆开处理。位置注意力管的是“哪里跟哪里有关系”通道注意力管的是“哪个特征跟哪个特征有关系”。如果你把它们揉在一个模块里就会得到一个四维的张量关系参数量和计算量完全失控而且梯度也不容易稳定。拆成两个分支后每个分支各自负责一种关系建模复杂度分别是O(N²)和O(C²)可控且互补。1.3 适合谁读谁可以跳过部分内容如果你做语义分割、实例分割、场景解析或者想给检测、跟踪模型加一个全局上下文模块这篇值得精读。它不依赖特别复杂的数学只要能看懂矩阵乘法就能理解全部推导。如果你已经熟悉Non-local网络或者自注意力机制那位置注意力模块对你来说基本没有理解门槛可以重点看通道注意力模块和两者的融合策略。如果只是想在业务场景里直接接入这个结构建议至少把两个分支的矩阵维度搞清楚不然后面debug的时候很容易摸不着头脑。2. 双注意力机制的核心模块拆解2.1 位置注意力模块像素之间的“隔空对话”位置注意力模块要做的事情可以用一句话概括让每个像素去询问所有其他像素“我们像不像”。像的话我多参考你的特征不像我少参考。这和自注意力机制是完全一致的思路。具体过程是这样的。假设输入特征图是A先通过两个1×1卷积分别生成B和C两个特征图目的是降维、减少后续矩阵乘法开销。然后把B和C拉平成二维矩阵形状从B×C×H×W变成B×C×NNH×W。接着做一次矩阵乘法B的转置和C相乘得到S矩阵形状是B×N×N。这个S矩阵里的每个元素就是位置j对位置i的相似度分数。为了让权重可解释、梯度更稳对S的每一行做softmax归一化。最后把value特征D也是由A经过一个1×1卷积得到的和S的转置相乘得到加权后的特征再还原成B×C×H×W乘一个可学习的缩放系数后和原始输入做残差相加。为什么要在softmax之后转置再和D相乘而不是直接乘这其实对应矩阵乘法的顺序。S矩阵的第i行表示位置i去匹配所有位置j得到的权重所以输出特征的第i个位置应该聚合D中所有位置的特征权重就是S的第i行。写成矩阵运算就是S与D转置相乘之后整个结果需要再做一次转置才能得到B×C×N。很多第一次复现的同学在这里绕晕建议亲手用一个小矩阵推一遍比自己琢磨半小时更快。2.2 通道注意力模块特征通道之间的自我校准通道注意力模块的出发点和位置注意力不一样。位置注意力关心的是空间关系通道注意力关心的是语义关系。一个直观的观察是如果某个通道专门负责检测轮子另一个通道负责检测车身那么这两个通道在面对车辆图片时应该会同时激活。通道注意力就是想显式地把这种“共现关系”建模出来。做法也很直接把输入特征A直接拉平成C×N的矩阵不做任何卷积变换。然后让A和A的转置相乘得到C×C的通道关系矩阵再对每一行做softmax。最后把这个通道关系矩阵和A相乘输出形状是C×N还原回空间维度后残差相加。也就是说每个通道的新特征是原始所有通道特征的加权和权重取决于通道之间的相似度。这里需要注意一个和SE模块的区别。SE模块是先做全局平均池化再用两个全连接层压缩出每个通道的标量权重本质上是对通道做“重要性重标定”不管空间关系。通道注意力模块则是直接计算通道两两之间的完整关系矩阵信息量更大但计算量也更大。实际使用时有人觉得SE更轻量有人觉得通道注意力的全局建模能力更强具体选哪个要看任务和资源。2.3 位置分支和通道分支为什么要并联再相加这是我在读的时候反复想过的一个问题为什么两个分支要做并联而不是串联试想如果做串联第一个分支的输出会改变特征分布第二个分支输入的是被第一个分支“污染”过的特征梯度信号也会链式传导得更深训练稳定性和收敛速度都更难控制。并联的结构让两个分支各自独立处理原始输入x各自产生一个残差最后相加。这样每个分支的梯度路径都相对独立训练起来稳定很多。相加之后有些实现会再接一个1×1卷积去融合这两个分支的信息有些实现直接相加就结束了。我的经验是如果backbone的信息量足够直接在两个分支输出相加后接分割头效果就不错如果任务比较复杂可以在相加后加一个轻量的卷积融合层让网络自己决定两个分支的权重而不是固定五五开。另外两个分支里那个可学习的缩放系数很重要。在原论文的实现里缩放系数初始化为0这意味着训练一开始模块退化为恒等映射然后再逐步放大注意力信号。这样设计的好处是模型在早期可以稳定地学习避免注意力矩阵一上来就给特征图带来剧烈扰动。3. 从公式到代码手把手实现Dual Attention模块3.1 特征图预处理reshape、转置与注意力矩阵的维度动手写代码之前一定要先把几个变量的形状理清楚。假设backbone输出的特征是X形状是B×C×H×W其中B是batch sizeC是通道数H和W是特征图的高宽。设置NH×W这是位置注意力模块的“token数量”也是整个模块最吃内存的地方。在实现中最常用的做法是把所有涉及空间维度的矩阵乘法都用批量矩阵乘法torch.bmm完成因为它专门处理B个独立的矩阵乘法效率比for循环高得多。需要注意的是经过view和permute之后张量的内存布局会发生变化最好加一个contiguous()否则后面调用bmm时可能报错。第一次写这个模块报错十有八九都是形状不匹配建议在forward里临时打印每个张量的shape调试。3.2 位置注意力模块的PyTorch实现与逐行说明下面是我自己整理的位置注意力模块实现基本按照论文里的过程来写加了一些注释方便看懂每一步在做什么。import torch import torch.nn as nn import torch.nn.functional as F class PositionAttentionModule(nn.Module): def __init__(self, in_channels, reduction8): super().__init__() # 用1x1卷积生成Q和K降低通道数以减少矩阵乘法计算量 self.query nn.Conv2d(in_channels, in_channels // reduction, 1) self.key nn.Conv2d(in_channels, in_channels // reduction, 1) # V保持原始通道数避免信息损失过大 self.value nn.Conv2d(in_channels, in_channels, 1) # 缩放系数初始化为0训练开始时模块等于恒等映射 self.gamma nn.Parameter(torch.zeros(1)) def forward(self, x): b, c, h, w x.size() n h * w # 生成Q、K、V并拉平成二维矩阵 q self.query(x).view(b, -1, n) # [b, c, n] k self.key(x).view(b, -1, n) # [b, c, n] v self.value(x).view(b, -1, n) # [b, c, n] # Q的转置与K相乘得到位置相似度矩阵 q q.permute(0, 2, 1).contiguous() # [b, n, c] k k.permute(0, 2, 1).contiguous() # [b, n, c] attn torch.bmm(q, k.permute(0, 2, 1)) # [b, n, n] attn F.softmax(attn, dim-1) # 对每一行做归一化 # 用注意力矩阵聚合value特征 out torch.bmm(v, attn.permute(0, 2, 1)) # [b, c, n] out out.view(b, c, h, w) out self.gamma * out x # 残差连接 return out这段代码里有一个关键细节位置注意力矩阵的softmax维度。我习惯对最后那个维度做归一化也就是让每一行之和为1。每一行代表一个源位置对所有目标位置的权重分布这样聚合出来的特征才有“当前位置参考全局”的含义。如果你把归一化维度搞反等价于每个位置都变成了“全局参考当前位置”实验效果会变得很怪而且训练时loss掉不下去。3.3 通道注意力模块的PyTorch实现与逐行说明通道注意力模块的实现更简单因为不需要额外的卷积变换直接对原始特征做自相关矩阵乘法就行。class ChannelAttentionModule(nn.Module): def __init__(self, in_channels): super().__init__() # 可学习缩放系数同样初始化为0 self.beta nn.Parameter(torch.zeros(1)) def forward(self, x): b, c, h, w x.size() n h * w # 拉平成 [b, c, n] a x.view(b, c, n) # 计算通道关系矩阵 attn torch.bmm(a, a.transpose(1, 2)) # [b, c, c] attn F.softmax(attn, dim-1) # 对每一行归一化 # 用通道关系矩阵加权原始特征 out torch.bmm(attn, a) # [b, c, n] out out.view(b, c, h, w) out self.beta * out x # 残差连接 return out通道注意力模块最需要注意的还是显存。C×C矩阵在通道数为256时是65536个元素看起来不大但在训练时需要保存梯度再乘上batch size开销会累积。如果你的backbone输出的通道数是512甚至1024这个矩阵乘法的开销会明显变大。我的建议是可以在进入通道注意力之前先加一个1×1卷积把通道数压缩到128或256做完注意力再还原这样能省不少显存效果通常损失很小。3.4 接入分割头融合模块的整体结构最后把两个模块组合成一个完整的Dual Attention Block然后接入语义分割头。class DualAttentionBlock(nn.Module): def __init__(self, in_channels): super().__init__() self.pam PositionAttentionModule(in_channels) self.cam ChannelAttentionModule(in_channels) # 融合层可选 self.fusion nn.Sequential( nn.Conv2d(in_channels * 2, in_channels, 1), nn.BatchNorm2d(in_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): pam_out self.pam(x) cam_out self.cam(x) # 两个分支的输出在通道维度上拼接再融合回原通道数 out torch.cat([pam_out, cam_out], dim1) return self.fusion(out)这一步有几种接法。原论文里两个分支的输出直接相加不额外拼接。我这里的写法是把两个分支输出拼起来再用1×1卷积融合。拼接的好处是网络能更灵活地决定两个分支各自贡献多少信息缺点是显存和参数量多了一点。如果你的模块是嵌入在某个轻量分割网络里直接相加更省事如果你在做一个精度优先的模型拼接融合效果通常会好一些。另外损失函数方面语义分割任务一般会同时加一个主损失和一个辅助损失。辅助损失通常接在backbone比较浅的位置帮助浅层特征也学到语义信息。Dual Attention Block通常只加在最后因为它的计算量不小如果每个stage都加训练会很慢。4. 复现路上的问题排查与调优经验4.1 显存爆炸N×N矩阵是最常见的坑位置注意力模块里有个N×N的相似度矩阵其中NH×W。假设backbone输出的特征图是64×64N就是4096N×N大约是1677万个元素每个元素float32占4字节单是一个注意力矩阵就需要大约64MB显存。看起来不多但训练时要保存中间梯度实际占用往往是这个数值的2到3倍。如果特征图变成128×128N是16384N×N就膨胀到2.68亿个元素单矩阵就要1GB以上这还没算batch size带来的倍增。所以我的建议是Dual Attention Block最好在分辨率比较低的特征图上使用比如原图1/8或者1/16下的输出。如果显存还是不够优先把位置注意力模块去掉只保留通道注意力模块效果通常不会掉太多。还有一个更省显存的办法是减少Q和K的通道数比如把in_channels除以8或除以16因为N×N矩阵的尺寸和通道数无关但Q、K计算和矩阵乘法本身的显存占用会变小。如果你的显存实在紧张又不想牺牲效果可以用梯度检查点技术在backward时重新计算前向的注意力矩阵而不是保存下来这样能大幅减少显存占用代价是训练时间变长。我个人的经验是先固定batch size等于1去跑通整个流程再逐渐加大batch比一上来就开大batch结果爆显存要省心得多。4.2 训练不收敛或者loss异常先检查这几个细节我复现时遇到的最常见问题是注意力矩阵里softmax的维度写错了。位置注意力应该在空间维度上归一化通道注意力应该在通道维度上归一化。如果写反loss会表现为前期下降一点点之后完全不动或者整个曲线变成一条平线。排查方法很简单把两个模块单独拿出来输入一个随机张量跑一下forward检查输出张量每个通道的数值范围是不是合理的。第二个常见问题是学习率。注意力模块和普通卷积不太一样它对学习率更敏感。刚开始训练时如果全局学习率设置得太大注意力矩阵的数值会被推到一个极端导致特征图被过度扰动。我建议先使用warmup策略前几个epoch让学习率从很小的值慢慢爬升等loss进入稳定下降区间之后再切换成poly或者cosine调度。在我的模拟项目X里同样的结构和数据加了warmup之后收敛速度比不加明显快最终指标也更高。第三个细节是BatchNorm。如果你在多个GPU上做同步训练默认的BatchNorm每个卡单独计算均值和方差这个对注意力模块影响不小。注意力矩阵本质上是在建模全局关系它对特征分布非常敏感每个卡上的统计量不一致会导致注意力权重失真。建议在多卡训练时改用SyncBN让batch统计量在所有卡上同步计算。4.3 加了模块效果反而下降先确认这两件事很多人复现注意力模块时遇到的最打击人的情况是加了这个模块之后指标反而比baseline还低。我遇到这种情况时会优先检查两件事。第一件事是数据量是否足够。注意力模块参数量不算大但它会把模型往“全局关系”这个方向带偏如果数据集本身比较小很容易过拟合到训练集的一些全局模式上。这时候可以尝试把缩放系数gamma和beta固定住或者降低它们的上界让模块弱化一些。另一个有效的做法是加更强的数据增强尤其是随机的尺度扰动和颜色扰动让模型不至于过拟合到固定的全局位置关系上。第二件事是模块放在模型里的位置是否合适。我在项目里测试过Dual Attention Block放在网络最底层那一层特征上时效果提升通常比较小因为底层特征本身还带着很多细节纹理全局关系反而可能把局部判别信息抹掉。放在靠近分割头的最后一层或者倒数第二层效果最好。如果你用的是多尺度特征的网络可以考虑把它加在融合后的特征图上让注意力模块去建模融合后的全局上下文。4.4 常见问题速查表问题现象可能原因解决建议显存溢出N×N注意力矩阵太大降低特征图分辨率、减小batch size、减少Q/K通道数、使用梯度检查点loss不下降或下降缓慢softmax维度写错、学习率过大单模块debug输出数值范围、加warmup、降低学习率训练初期loss剧烈震荡缩放系数初始值过大将gamma和beta初始化为0或设置很小的初始值多卡训练效果比单卡差BatchNorm统计量不一致换成SyncBN加了模块效果不如baseline模块位置不对、数据量不足尝试放到倒数第二层、加强数据增强、弱化缩放系数推理速度慢N×N矩阵乘法耗时长用更低分辨率特征图、尝试只保留通道注意力、或用线性注意力近似5. 读完之后的扩展实践与思考5.1 从分割到检测注意力结构的迁移价值读这篇paper的时候我一直在想这种双注意力结构能不能用到检测任务上。检测里有一个经典问题小目标和大目标同时存在时模型容易对小目标召回不足对遮挡目标分类错误。Dual Attention里的通道注意力分支可以看作是一种自适应的特征重加权机制它让某些语义相关的通道互相增强这对提升小目标的召回率很有帮助。而位置注意力分支则可以让被遮挡的目标借助周围环境信息推断出完整语义比如被挡住一半的行人也能被正确识别。我在某个目标检测项目里试过把通道注意力模块单独拿出来接到检测头的输入特征上效果有一定提升但位置注意力分支因为N×N矩阵开销太大当时没有直接用改成了局部窗口版本让每个位置只和周围一定范围内的位置计算相似度。这样既保留了全局建模的思想又把复杂度从N×N降到了N×k在工程上更可行。5.2 复杂度的进一步优化方向如果要做更大分辨率输入或者实时推理原始的N×N矩阵乘法还是太奢侈。现在业界很流行的做法是把注意力矩阵做线性近似比如用核函数展开或者低秩分解把N×N计算量降成线性级别。你完全可以保留Dual Attention的双分支结构只是把里面的相似度计算替换成线性版本。我自己的体会是先别急着换复杂结构先把原始版本跑通复现出合理的精度提升再去尝试各种改进。很多初学者一上来就直接上改进版结果出了问题根本不知道是原始结构的问题还是改进带来的问题。在模拟项目X里我一开始也是直接改进后来老老实实把原始Dual Attention复现了一遍才真正理解了每个分支的作用后面的改进才变得有迹可循。对我来说读这篇paper最大的收获不是那套矩阵公式而是一种设计思路两件独立且互补的事情分开建模再融合往往比硬塞进同一个模块里更稳也更好调试。后续如果想继续深挖我打算把Dual Attention和Non-local做一个系统的对比实验看看在相同计算量下哪个结构对语义分割更友好。最后给一个很实际的小建议复现这类注意力模块时一定先把随机种子和损失固定住用一个小的验证集反复对比加模块前后的差距这样才能真正看出注意力结构到底帮了多少。这一点我踩过好几次坑写出来希望大家别再走弯路。