最近在读语义分割方向的论文,正好看到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×N,N=H×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 size,C是通道数,H和W是特征图的高宽。设置N=H×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, reduction=8): 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(inplace=True) ) def forward(self, x): pam_out = self.pam(x) cam_out = self.cam(x) # 两个分支的输出在通道维度上拼接,再融合回原通道数 out = torch.cat([pam_out, cam_out], dim=1) return self.fusion(out)这一步有几种接法。原论文里两个分支的输出直接相加,不额外拼接。我这里的写法是把两个分支输出拼起来,再用1×1卷积融合。拼接的好处是网络能更灵活地决定两个分支各自贡献多少信息,缺点是显存和参数量多了一点。如果你的模块是嵌入在某个轻量分割网络里,直接相加更省事;如果你在做一个精度优先的模型,拼接融合效果通常会好一些。
另外,损失函数方面,语义分割任务一般会同时加一个主损失和一个辅助损失。辅助损失通常接在backbone比较浅的位置,帮助浅层特征也学到语义信息。Dual Attention Block通常只加在最后,因为它的计算量不小,如果每个stage都加,训练会很慢。
4. 复现路上的问题排查与调优经验
4.1 显存爆炸:N×N矩阵是最常见的坑
位置注意力模块里有个N×N的相似度矩阵,其中N=H×W。假设backbone输出的特征图是64×64,N就是4096,N×N大约是1677万个元素,每个元素float32占4字节,单是一个注意力矩阵就需要大约64MB显存。看起来不多,但训练时要保存中间梯度,实际占用往往是这个数值的2到3倍。如果特征图变成128×128,N是16384,N×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做一个系统的对比实验,看看在相同计算量下哪个结构对语义分割更友好。最后给一个很实际的小建议:复现这类注意力模块时,一定先把随机种子和损失固定住,用一个小的验证集反复对比加模块前后的差距,这样才能真正看出注意力结构到底帮了多少。这一点我踩过好几次坑,写出来希望大家别再走弯路。