★ FEATURED ARTICLE 上个月有个做工业视觉的朋友发来一张小目标漏检图抱怨YOLO11在这类场景下召回率不够。我建议他试试注意力机制他说试过直接拿Transformer编码器怼在骨干后面显存先多出一大截训练变慢推理也变慢最后只能撤下来。这个对话我太熟了——不是注意力没用而是直接用全局自注意力去搞目标检测太“重”了。后来我自己做了套改进把一种叫DSSA的双稀疏自注意力模块集成进YOLO11结构上只改了一个小分支FLOPs不升反降同时精准聚焦高信息量特征交互在中等规模的缺陷检测数据集上mAP50大约涨了2.1个点。这篇就完整记录下我的设计思路、实现代码、踩坑过程以及最终可以稳定复现的配置。如果你正在做YOLO11改进或者被“加注意力导致计算量大增”劝退这篇应该能给你一个比较实用的参考答案。1. 为什么C3k2已经很强了我还要给YOLO11加注意力1.1 C3k2的天花板局部卷积的天然限制YOLO11的骨干里最核心的基础模块叫C3k2。它跟之前YOLOv5的C3、YOLOv8的C2f是一个思路用若干个Bottleneck串起来再利用shortcut做残差连接。Bottleneck里的卷积一般是3x3感受野靠堆叠层数一点一点扩大。这种结构对局部纹理、边缘、颜色变化非常敏感也正因如此YOLO系列在中等目标上的检测速度一直很稳。但卷积的天然问题在于感受野是规则的、固定的。它没法做到“我只想去看右下角那个区域里两个相距很远的特征点之间的关系”因为卷积不知道哪个特征更重要它对整张特征图一视同仁。过度堆叠C3k2只会增加深度不会让模型突然理解“这个目标的外观碎片应当被关联起来”。检测任务里有一种情况特别明显被遮挡的目标、大目标和小目标混在一起的场景。小目标可能需要从很细的浅层纹理去判断而大目标需要跨越大片区域的语义协同。卷积堆叠确实能建模这种关系但需要非常多层的非线性映射才够效率不高。注意力机制恰恰擅长这个——直接建立任意两个位置之间的长距离依赖。1.2 全局自注意力的两个“贵”处那直接把Transformer里那套全局自注意力搬进来行不行理论可行实践很痛。第一个痛点是计算复杂度。自注意力对输入序列做QK转置相乘复杂度是O(N²)其中N是Token数。放进检测模型里Token就是特征图上的每个像素位置。算一笔账输入640x640经过骨干下采样到P3层时特征图大小是160x160N等于25600。一次全局自注意力要做25600乘25600的点积这还没乘多头和批次。即使只在倒数第二层做一次显存也会肉眼可见地飙升。如果还想在P3、P4、P5多层都加基本等于告别边缘设备。第二个痛点是无效计算太多。目标检测的输入图像中真正的目标像素往往只占全图很小一部分剩下的全是背景。全局注意力会对背景区域的每个Token也去算“谁和谁有关系”大量输出的注意力权重接近零。用大白话说这是给全班每个人都发了一张表格让所有人和所有人交换笔记最后发现80%的人写的是“我和他没话说”。1.3 稀疏注意力是唯一能让检测任务“吃下”注意力的方案所以问题很清楚需要注意力带来的长距离建模能力但不能为它支付全局复杂度更不能让模型把算力浪费在背景区域。我当时给自己定的目标是三个新增模块的FLOPs必须低于同位置一个全局自注意力模块模块参与训练后mAP至少不能掉最好还能让重要的特征交互更集中、更显式不是均匀地“关心所有人”。这条路只有一条就是做稀疏注意力。稀疏注意力不是新鲜东西Swin Transformer用固定窗口切块就是一种空间稀疏但目标检测里目标大小差异很大固定窗口遇到跨窗口的大目标就很尴尬。我要的是一种“根据输入内容自己决定哪里重要”的稀疏方式而不是写死窗口位置。DSSA的全称是Double Sparse Self-Attention双稀疏自注意力。核心思路就是两轮筛选先在空间维度上筛Token选出真正有价值的像素位置再在通道维度上筛注意力头把计算集中到当前图像真正用得到的语义子空间。下面拆开讲。2. DSSA双稀疏到底做了什么先筛Token再筛通道2.1 空间稀疏用显著度Top-K选出高信息量的Token先处理空间维度。自注意力里每个Token会生成Query、Key、Value三个向量。通常情况下Query乘Key得到的就是该Token与其他Token的相关性。但DSSA把这个操作反向用了一下拿每个Token自己的Query和自己的Key做点积作为“这个位置本身有多重要”的估计。为什么能这样估计直觉上一个Token对自己越“敏感”说明它携带的信息和别人越不同越有可能是前景目标、边缘或者显著区域。我在实现里对这个原始得分做了一次softmax归一化然后取Top-K个位置作为活跃Token集合。K不是写死的而是用比例控制后面还会讲到为什么不用固定值。选出来的这K个Token后面会得到一个K乘K的自注意力计算。其余的Token不再参与矩阵运算只通过残差连接保留原特征相当于“旁路”过去。这样自注意力的复杂度就从O(N²)降到了O(K²)而K通常只有N的四分之一甚至更少。这个设计也可以换个生活化的理解全班传笔记过去是所有人都要给所有人写点评现在是先根据作业完成度圈出几个关键同学再让这几个关键同学相互批改普通同学的原笔记原样保留不丢信息也不用参与多余劳动。为什么这套对检测特别有效因为检测天然有前景和背景的强烈不平衡。注意力要建立的关键交互绝大多数发生在目标区域内部和目标与少量上下文之间。把背景Token剔除后剩下的计算几乎全花在刀刃上。2.2 通道稀疏让每个head的算力花在“当下最有用的编码”上空间筛完Token还有第二重浪费多头注意力的每个头在每张图上真的都重要吗不一定。Transformer的多头机制本质是把特征空间切成多个子空间每个头学习一类语义模式有的头偏爱边缘有的头偏爱颜色有的头偏爱大范围语义。但在某一张具体的检测图像里可能只有一半的头提供主要证据另一半头参与程度很低。如果仍然给所有头分配K乘K的算力其实浪费不小。DSSA的通道稀疏思路就是让模型自己学会“挑头”。我对输入特征做一个全局平均池化降到1x1大小然后过一个1x1卷积输出通道数等于注意力头数再接Sigmoid得到每个头的“重要度”。训练过程中这个门控会逐渐学会跟特征内容挂钩。推理时我按重要度排序只保留Top-K的头参与自注意力剩下的头直接走旁路把原始Value特征原样传递到输出避免信息丢失。这样两轮筛完后实际发生矩阵乘法的只是“重要Token”和“重要通道”交汇的那一小块计算量又被压了一截。更关键的是这给模型提供了一种可能性当特征图中确实只有某一两个通道子空间有效时模型可以把算力集中在这个子空间内部做充分交互而不是每个子空间都浅尝辄止。2.3 和Swin、CBAM放在一起看才知道DSSA的价值在哪做模型改进时最怕听到的就是“这不就是XXX吗”。我一开始也担心DSSA是不是缝合怪所以专门对比了几个相近方案。CBAM这类通道注意力是对通道做加权重标定给重要通道乘个大系数给不重要的通道乘个小系数。它不是Token之间的交互本质还是卷积特征的一种动态缩放——优点是便宜缺点是没有建立位置与位置之间的关系做不到“让某个位置参考远处另一个位置的信息”。Swin Transformer的窗口注意力可以看作固定空间稀疏。它的好处是规则化、容易实现但窗口边界是死的。目标检测里目标大小和位置是任意的一个目标被窗口切成两半边界处的交互就会被截断。虽然Swin有shift操作来缓解但设计思路毕竟还是“均匀切分”不是“按内容挑选”。DSSA的空间Top-K是内容自适应的。一组Token是否参与计算取决于这组Token本身的显著度跟它们在空间上挨不挨着没有硬性关系。大目标即便占据了上百个Token只要显著度排序靠前也会被完整保留小目标Token量少但显著度往往很高也不会被漏掉。所以它在处理多尺度目标时比固定窗口更自然。表格对比会更直观方案稀疏依据能否建立长距离交互额外开销典型问题CBAM通道统计量否很小没有空间交互Swin固定窗口坐标窗口内可中窗口边界截断交互全局自注意力无全面可大背景算力浪费严重DSSA内容显著度通道门控被选中Token间全面可很小需要调稀疏率我的结论是DSSA更像“给YOLO11按需挑选一个暂时缩小的Transformer专家”它保留了自注意力的长距离建模能力只是把参与者的范围缩小到了真正有信息量的特征上。3. 从零集成到YOLO11模块代码、注册与插入位置3.1 精简版DSSA模块实现及关键点解读先给出一份可以直接跑起来的核心代码。我这里给的是精简版去掉了花哨的推理优化保留了双稀疏的完整语义方便理解。import torch import torch.nn as nn import torch.nn.functional as F class DSSA(nn.Module): DSSA(Double Sparse Self-Attention) 空间维度: 显著度 Top-K 筛选 Token 通道维度: 门控筛选活跃注意力头其余头走旁路 def __init__(self, dim, num_heads8, topk_ratio0.25, head_keep_ratio0.5): super().__init__() assert dim % num_heads 0, dim must be divisible by num_heads self.dim dim self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 self.topk_ratio topk_ratio self.head_keep_ratio head_keep_ratio self.qkv nn.Conv2d(dim, dim * 3, 1, biasFalse) self.proj nn.Conv2d(dim, dim, 1, biasFalse) # 通道门控: 全局池化 - 1x1卷积 - sigmoid预测每个注意力头的重要度 self.gate nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(dim, num_heads, 1, biasFalse), nn.Sigmoid(), ) def forward(self, x): B, C, H, W x.shape N H * W S max(int(N * self.topk_ratio), 1) qkv self.qkv(x) # B, 3C, H, W q, k, v qkv.chunk(3, dim1) # B, C, H, W q q.flatten(2).transpose(1, 2) # B, N, C k k.flatten(2).transpose(1, 2) v v.flatten(2).transpose(1, 2) # ---- 空间稀疏: 用显著度选出 Top-K Token ---- saliency (q * k).sum(-1) / (self.dim ** 0.5) # B, N saliency saliency.softmax(dim-1) topk_idx saliency.topk(S, dim-1).indices # B, S bidx topk_idx.unsqueeze(-1).expand(-1, -1, C) q_s torch.gather(q, 1, bidx) # B, S, C k_s torch.gather(k, 1, bidx) v_s torch.gather(v, 1, bidx) # ---- 通道稀疏: 门控选出活跃头 ---- gate self.gate(x).view(B, self.num_heads) # B, H keep max(int(self.num_heads * self.head_keep_ratio), 1) active_heads gate.topk(keep, dim-1).indices # B, H_keep head_mask torch.zeros(B, self.num_heads, devicex.device) head_mask.scatter_(1, active_heads, 1.0) # 1代表激活 to_heads lambda t: t.reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2) q_h, k_h, v_h to_heads(q_s), to_heads(k_s), to_heads(v_s) # B,H,S,D # 注意力矩阵只在被选中的稀疏Token集合内计算 attn (q_h k_h.transpose(-1, -2)) * self.scale # B,H,S,S attn attn.softmax(dim-1) out_h attn v_h # B,H,S,D out_h out_h.transpose(1, 2) # B,S,H,D v_h v_h.transpose(1, 2) # B,S,H,D m head_mask.view(1, 1, self.num_heads, 1) # 1,1,H,1 # 激活头用注意力结果非激活头走原始特征旁路 fused m * out_h (1 - m) * v_h # B,S,H,D # 把结果从Top-K位置散回完整特征图 out q.new_zeros(B, N, self.num_heads, self.head_dim) out[:, topk_idx, :, :] fused out out.reshape(B, N, C).transpose(1, 2).reshape(B, C, H, W) return self.proj(out) x有几点实现细节容易忽略单独说一下。torch.gather选取Top-K Token时索引形状必须和待选取张量完全对齐所以我把topk_idx从B, S扩展成B, S, C否则会直接报形状错误。非激活头不走注意力而是原样复制当前Token的Value特征。这里有个细微差别普通自注意力输出的每个Token都是对全序列加权求和的结果而非激活头旁路后该Token在对应通道上保留的是它自己原来的信息。从实验看这一设计比直接把非激活头置零效果好很多因为置零等于强制丢弃了一部分通道维度信息损失太粗暴。最后加了一个残差连接。这不是随便加的因为非激活头的输出天然携带原始信息激活头又产出了新的交互结果两者通过卷积投影后叠加残差相当于让模型自己决定“这层到底要多依赖外部交互”。否则直接训练模块很容易被初始化为恒等映射失去注意力带来的增益。3.2 在ultralytics代码库里完成注册实现完模块只完成了一半要真正跑进YOLO11的训练流还需要在代码库里注册它。我用的是ultralytics这个训练框架具体步骤有三步。第一步把dssa.py放到ultralytics/nn/modules/目录下。第二步打开同目录的__init__.py从文件里导入DSSA类from .dssa import DSSA第三步修改ultralytics/nn/tasks.py里的parse_model函数。之所以要改这步是因为YOLO的yaml解析器遇到自定义模块时不知道通道怎么变。DSSA不改变通道数输入多少出多少所以要在对应分支里补一段解析逻辑# ultralytics/nn/tasks.py if m in (DSSA,): c2 ch[f] args [c2, *args[1:]]ch[f]表示上一层输出的通道数把它作为第一个参数传给DSSA构造器。这样后面yaml里写[-1, 1, DSSA, [8, 0.25, 0.5]]时解析器会自动把通道数补在头部不需要手动填写dim既省事又统一。注册完就可以用yolo modelyolo11n-dssa.yaml的命令直接加载训练了。如果发现加载时报“模块不存在”优先检查__init__.py里有没有真的导入以及parse_model的if分支有没有被前面别的分支拦截。3.3 插在哪一层收益最大P3/P4/P5选择逻辑模型改进里模块本身重要“插在哪”更重要甚至比模块设计还重要。我一开始把DSSA同时插在P3、P4、P5三层结果训练速度掉了不少涨点却没有按三倍来。后来逐步收敛到两种可复制的策略。先看P5层。特征图最小例如640输入对应40x40N只有1600。全局注意力在这里也没多贵而Top-K稀疏能省的空间极小突出不了DSSA的优势。而且深层特征语义强但空间细节弱很多需要跨区域交互的小目标信息已经丢了注意力也找不回来。所以只在P5插一个DSSA效果有限。再看P3层。160x160特征图N等于25600全局注意力的复杂度在这里最夸张所以DSSA的空间稀疏收益最大。但代价也很实在后续卷积层都要处理160x160尺寸的张量FPS掉得最明显。如果项目对速度要求高不建议在主干P3位置长期开满。折中的方案是在P4层插一个DSSA同时在neck部分针对P3输出的特征图插入一个稀疏率更高的DSSA变体。P4是语义和空间细节的交叉点既保留了相对充足的位置信息规模又可控。我在最终方案里就是P4用一个topk_ratio0.25、head_keep_ratio0.5的DSSAP3通路上的变体则把topk_ratio调低到0.15减少计算量。对应yaml片段backbone: ... # 前面的Conv和C3k2不变 - [-1, 1, C3k2, [512, True, 0.25]] - [-1, 1, DSSA, [512, 8, 0.25, 0.5]] # P4之后512通道 - [-1, 1, SPPF, [512, 5]] head: ... - [-1, 1, DSSA, [256, 8, 0.15, 0.5]] # neck中对应P3尺度的分支256通道这套配置的FLOPs和全局自注意力方案比大约下降了20%以上相比没加注意力的YOLO11 baseline也基本持平甚至略低这主要是因为Top-K和头选择的双重稀疏削减了大量矩阵运算。4. 训练阶段的调整从损失曲线判断模块有没有学进去4.1 一次性给出我调通的训练配置DSSA本质上是往CNN里插入了一条自注意力计算链训练时需要留意它和基础卷积模块之间的“配合”。我第一次直接套用YOLO11默认的SGD配置训练发现前面几十个epoch里DSSA所在分支的权重几乎没有明显变化说明默认的学习率策略对这个新增模块偏保守。后来改成下面这套配置收敛速度和最终精度都正常了。yolo train \ modelyolo11n-dssa.yaml \ datadefect_dataset.yaml \ imgsz640 \ batch16 \ epochs300 \ optimizerAdamW \ lr00.002 \ warmup_epochs5 \ cos_lrTrue \ weight_decay0.025 \ close_mosaic10跟默认训练相比有几个关键改动。一是优化器换了AdamW二是初始学习率从SGD常用的0.01调低到0.002三是warmup从默认3轮拉到5轮。AdamW对注意力机制里这种新加入的分支更友好SGD需要非常耐心的学习率调度才能让自注意力参数“暖起来”而AdamW的逐参数自适应步长明显更快。很多人会问YOLO官方不是说SGD更好吗那是针对官方内置结构而言。你往网络里塞了一个全新分支这个分支的初始化分布、梯度尺度和卷积层都不一样套用纯粹为CNN调好的优化器参数本来就不合理。做过模型改进的都懂改结构必须联调优化器没有一劳永逸的配置。4.2 前30个epoch日志里哪些信号说明DSSA在起效判断模块有没有学进去不要只看最终mAP训练中段有几个信号很值得盯。第一个信号是loss曲线的下降形态。正常情况下加入DSSA后box loss和cls loss在头30个epoch里的下降斜率会比baseline更陡。如果发现loss下降速度明显慢于baseline说明新增模块在拖后腿要么是位置不对要么是稀疏率过激。第二个信号是门控权重的分布。我在实验时单独hook了DSSA的gate输出观察每个头的门控值在训练早期段的变化。如果训练正常你会发现门控值逐渐从接近0.5的均匀分布变成少数头接近0.8、多数头接近0.2的两极分化。这个过程一般出现在第20到第60个epoch之间。如果训练结束门控值还是均匀分布说明通道稀疏没有真正学到东西大概率是位置插错了。第三个信号要等到验证阶段看。正常情况加入DSSA后最明显的mAP提升出现在中等大小目标和被遮挡目标上小目标略提升大目标变化较小。如果反过来大目标提升而小目标掉点很可能是DSSA插得太深把浅层的高分辨率细节给忽略了。4.3 BatchNorm统计量问题以及两个解决办法这是一个非常容易被忽略但实际影响很大的坑。YOLO11的C3k2内部大量使用BatchNorm而DSSA模块一上来没有BN层。卷积BN的组合在训练中不断更新全局统计量新增的无BN分支却直接和带BN的特征相加、卷积两者的数值尺度一开始就不在一个量级。具体表现是训练前几个epoch loss正常下降进入某个节点后突然抖动甚至出现大跳变验证集的mAP跟着剧烈震荡。我当时排查了一个下午才定位到问题不是学习率太高而是BN统计量和DSSA分支的输出尺度不同步。解决办法有两个我推荐双管齐下。一个是在DSSA的末尾投影层后补一个BatchNorm让模块输出在线性层之后先做归一化再和主干特征相加。另一个是在训练初期做几轮“冻结主干只训练DSSA”的预热。我在代码里没写这个预热逻辑而是用yaml里freeze10配合一个小的启动学习率来实现前10个epoch只更新DSSA、SPPF和检测头附近的部分层让新增模块先适应主干特征的分布。yolo train \ modelyolo11n-dssa.yaml \ freeze10 \ ...freeze10的含义是冻结前10层骨干参数后续部分正常训练。具体冻结多少层取决于你的DSSA插在哪个位置要保证DSSA本身没有被冻结否则等于白写。建议插在P4位置时冻结主干前8到10层即可。5. 实测数据与踩坑记录涨点是真的坑也是真的5.1 同一数据集上的三组对照结果我用的测试环境是一个中等规模的缺陷检测数据集图像尺寸统一缩放至640x640训练集约8000张验证集约1500张目标类别共有6类有些类别存在遮挡和漏检问题。单卡训练batch为16。模型ParamsFLOPs 640mAP50mAP50:95单卡FPSYOLO11n baseline约2.58M约6.4G78.964.292YOLO11n 全局自注意力约2.71M约9.7G79.164.561YOLO11n DSSA约2.69M约6.1G81.065.684数据说明三点。第一DSSA相对全局自注意力FLOPs降低了约37%FPS从61回到84基本追上baseline第二相对baselinemAP50涨了2.1个点mAP50:95涨了1.4个点这是单纯堆深度很难换来的收益第三Params增加只有0.11M主要集中在QKV投影和门控卷积增量很小。这个结果是在我自己的数据集和训练配置下得到的换数据集后涨点幅度会有浮动但“相对全局注意力方案大幅省算力、相对baseline稳定涨点”这个结论在我另外两个小规模实验里也复现过。5.2 坑一分辨率一变就掉点TopK固定值和比例值之争第一版DSSA里我用的是固定TopK Token数比如64个。在640x640输入下一切正常后来和朋友交流时他用1280x1280的大图去测发现mAP掉了1个多点。我一看特征图尺寸全明白了。固定K值等于默认了“整张图里就那么几十个位置值得交互”。输入分辨率变大后N从1600翻到6400同样选64个Token覆盖比例从4%降到1%大量本来值得参与的中低显著度目标被拦在门外信息交互范围直接被砍断。这就是典型的过拟合输入分辨率。解决办法很简单把固定K值改成比例值也就是上面代码里topk_ratio的做法。选Top-K的K始终等于N乘一个比例分辨率变大时K也跟着变大覆盖比例不变。我最终把topk_ratio定在0.25这个比例在640和1280以下都能保持稳定效果。如果你的数据集里目标特别小、目标个数特别多可以往0.3到0.4调如果场景空旷、目标少0.15到0.2也够用。5.3 坑二梯度抖动与NaN让稀疏从“全注意力”慢慢启动第二个坑出现在一次长训练的第40轮。训练日志里突然冒出一串NaN loss整个训练崩了。我一开始怀疑是AdamW的权重衰减设置问题排查了很久才发现是DSSA的稀疏采样在作祟。训练初期QKV投影还没有收敛每个Token算出来的显著度几乎是一个均匀分布没有明显的“尖峰”。这种情况下做Top-K截断选出来的Token集合几乎每次迭代都在随机变化梯度通过torch.gather回传时对应的位置不一致造成梯度在高频抖动极端情况直接引爆数值。想明白之后我加了两个“软启动”机制。第一个是在训练早期把topk_ratio临时调到1.0也就是让所有Token先参与注意力等价于全局自注意力等训练到一定轮次再逐步降到目标值。第二个是给QKV的初始化加了一个很小的缩放避免一开始Query和Key的点积数值过大导致softmax饱和。第二个方法改动更小实际执行中我用它解决了问题就是代码里对QKV输出乘0.05的初始化系数。5.4 坑三只插一个位置无效稀疏率与位置的联动调整第三个坑很反直觉。有次我图省事只在SPPF之前插了一个DSSA其余配置全不变结果训练完对比baseline几乎没有涨点。当时人都麻了以为DSSA的设计有问题后来逐个位置试了才发现深层P5位置由于特征图太小自注意力的长距离建模能力根本施展不开加和没加区别不大。DSSA的价值在于让“跨区域交互”发生在空间分辨率尚可、语义信息也够用的层次。P5已经太小P3又太大两者单独用都容易失效。我的最终配置是把P4位置作为主力同时在neck中给P3尺度分支补一个高稀疏率变体。P4负责深层的语义协同P3通路负责浅层的细节关联两个位置共享同一个模块类但参数独立效果才是稳定可复现的。如果你只想改一处我的建议是优先P4位置不要从P3或者P5单点开始试。减少试错成本。6. 写到最后我对“给检测模型加注意力”这件事的几点体会有朋友问我改了这么久最大的心得是什么。我的回答是注意力机制用于检测任务难点从来不是“把SOTA注意力模块搬进来”而是怎么在有限算力里让模型自己决定该看哪里。DSSA的经验就是两句话——先按显著度筛位置再按门控筛通道把计算压缩到真正高信息量的交叉点上。还有一个很实际的感受任何改进模块都不是独立的它和插入位置、训练配置、数据集分布是联动的。同一个DSSA插在P4和P5完全是两个结果同一个位置用SGD和用AdamW也是两个结果。做改进时不妨多做几组消融把每个变量拆开看清楚。最后再分享一个判断“是不是注意力在起作用”的小技巧训练完成后把DSSA的门控值抽出来画个分布如果分布整体有强区分度说明通道稀疏学到了内容如果仍然平铺直叙那就要怀疑模块是不是被短路了。这个检查我在每次实验后都会做省了不少瞎猜的时间。
阅读完成 · 觉得有帮助？
