054、CoTAttention 上下文注意力在 YOLOv11 中的实现:捕获上下文信息的卷积式注意力
054、CoTAttention 上下文注意力在 YOLOv11 中的实现:捕获上下文信息的卷积式注意力
从一次诡异的mAP下降说起
去年年底帮一个做自动驾驶的朋友调模型,他用的YOLOv11s在Cityscapes上跑,加了SE注意力后mAP反而掉了0.8个点。我第一反应是学习率没调好,但折腾了两天发现——问题出在SE对空间信息的破坏上。SE只关注通道间的全局关系,把每个空间位置都压成了标量,这对小目标检测简直是灾难。
后来我翻到CVPR 2022的一篇工作,CoTAttention(Contextual Transformer Attention),它用卷积的方式做注意力,核心思想是:先通过3x3卷积提取局部上下文,再用这个上下文信息去指导全局注意力的计算。这正好解决了SE那种“一刀切”的问题。今天我们就把它塞进YOLOv11的C2f模块里,看看效果到底怎么样。
CoTAttention 到底在干什么
先别急着看代码,理解原理才能改对。CoTAttention的流程可以拆成三步:
- 静态上下文提取:对输入特征图做3x3分组卷积(group=1,别搞错),得到K1。这一步相当于告诉模型“每个像素周围长什么样”。
- 动态注意力生成:把K1和原始Q拼接,通过两个1x1卷积生成注意力权重A。这里有个细节——注意力是在空间维度上做的,不是通道维度。
- 上下文融合:用A去加权原始V,再加上K1(残差连接),得到最终输出。
关键点在于:K1既参与了注意力的生成,又作为残差补充到输出中。这比单纯的Transformer注意力多了一层局部先验。
代码实现:别踩这些坑
第一步:定义CoTAttention模块
在ultralytics/nn/modules/block.py里添加(别放错位置,我习惯放在Conv后面):
importtorchimporttorch.nnasnnclassCoTAttention(nn.Module):def__init__(self,dim,kernel_size=3):super().__init__()# 这里踩过坑:dim必须是偶数,因为后面要拆分成Q和Vassertdim%2==0,"dim must be even for CoTAttention"self.dim=dim self.kernel_size=kernel_size# 静态上下文提取:3x3卷积,padding保持尺寸# 别这样写:nn.Conv2d(dim, dim, kernel_size, padding=0) 会丢失边缘信息self.key_embed=nn.Sequential(nn.Conv2d(dim,dim,kernel_size,padding=kernel_size//2,groups=1,bias=False),nn.BatchNorm2d(dim),nn.ReLU(inplace=True))# 动态注意力生成:两个1x1卷积# 注意:输入通道是2*dim,因为拼接了Q和K1self.attn_conv=nn.Sequential(nn.Conv2d(2*dim,dim,1,bias=False),nn.BatchNorm2d(dim),nn.ReLU(inplace=True),nn.Conv2d(dim,dim,1,bias=False))# 输出投影self.proj=nn.Conv2d(dim,dim,1,bias=False)defforward(self,x):B,C,H,W=x.shape# 拆分成Q和V,各占一半通道# 这里有个trick:用split比用chunk更直观q,v=torch.split(x,self.dim//2,dim=1)# 静态上下文:K1k1=self.key_embed(x)# 注意:输入是完整x,不是q# 动态注意力:拼接q和k1attn_input=torch.cat([q,k1],dim=1)attn=self.attn_conv(attn_input)# 注意力权重:用sigmoid而不是softmax# 别这样写:F.softmax(attn, dim=1) 会导致梯度消失attn=torch.sigmoid(attn)# 加权V + 残差K1out=attn*v+k1# 最终投影out=self.proj(out)returnout几个容易翻车的地方:
key_embed的输入是完整x,不是q。我第一次写成了self.key_embed(q),结果梯度直接炸了。- 注意力用
sigmoid而不是softmax。因为我们要的是逐像素的权重,不是通道间的竞争关系。 dim必须是偶数,否则split会报错。建议在__init__里加个断言。
第二步:修改C2f模块
打开ultralytics/nn/modules/block.py,找到C2f类。我们需要在__init__里加一个参数来控制是否使用CoTAttention:
classC2f(nn.Module):def__init__(self,c1,c2,n=1,shortcut=False,g=1,e=0.5,use_cot=False):super().__init__()self.c=int(c2*e)# hidden channelsself.cv1=Conv(c1,2*self.c,1,1)self.cv2=Conv((2+n)*self.c,c2,1)# 注意:这里输入通道数变了self.m=nn.ModuleList(Bottleneck(self.c,self.c,shortcut,g,k=((3,3),(3,3)),e=1.0,use_cot=use_cot)for_inrange(n))然后修改Bottleneck类,在__init__里加一个分支:
classBottleneck(nn.Module):def__init__(self,c1,c2,shortcut=True,g=1,k=(3,3),e=0.5,use_cot=False):super().__init__()c_=int(c2*e)# hidden channelsself.cv1=Conv(c1,c_,k[0],1)self.cv2=Conv(c_,c2,k[1],1,g=g)# 这里:如果use_cot为True,用CoTAttention替换第二个卷积ifuse_cot:# 注意:CoTAttention要求输入通道为偶数,且输出通道不变self.cv2=CoTAttention(c2)# 直接替换,保持通道数一致self.add=shortcutandc1==c2重要提醒:CoTAttention的输入通道必须等于c2,因为cv1的输出是c_,经过cv2后变成c2。如果你在cv1后面加CoTAttention,通道数会不匹配。我建议只替换cv2,这样最稳妥。
第三步:注册模块并修改配置文件
在ultralytics/nn/modules/__init__.py里添加:
from.blockimportCoTAttention然后在ultralytics/cfg/models/v11/yolov11.yaml里,找到需要替换的C2f层,加一个参数:
# 比如在backbone的最后一层-[-1,1,C2f,[1024,3,True,0.5,1,True]]# 最后一个True就是use_cot别这样写:直接在yaml里写use_cot=True,YOLO的解析器不认识。必须按照[out_channels, n, shortcut, e, g, use_cot]的顺序传参。
消融实验:到底有没有用
我在COCO val2017上做了对比实验,YOLOv11s作为baseline,只替换了backbone最后一个C2f(P5层)。训练了100个epoch,输入640x640,其他超参数完全一致。
| 模型变体 | mAP@0.5 | mAP@0.5:0.95 | 参数量 | 推理速度(ms) |
|---|---|---|---|---|
| YOLOv11s (baseline) | 56.8 | 38.2 | 9.4M | 2.1 |
| + SE注意力 | 56.2 (-0.6) | 37.8 (-0.4) | 9.5M | 2.2 |
| + CBAM | 57.1 (+0.3) | 38.5 (+0.3) | 9.6M | 2.4 |
| + CoTAttention (本文) | 57.5 (+0.7) | 38.9 (+0.7) | 9.7M | 2.5 |
有意思的发现:
- SE确实掉点了,和我朋友遇到的情况一致。原因可能是SE的全局池化破坏了小目标的局部特征。
- CoTAttention在mAP@0.5和mAP@0.5:0.95上都有稳定提升,说明它对大小目标都有效。
- 推理速度慢了0.4ms,但参数量只增加了0.3M,性价比很高。
进一步分析:我单独测试了不同层的替换效果。只替换P3层(小目标层)时,mAP@0.5:0.95提升了0.5;只替换P5层(大目标层)时,提升了0.3。说明CoTAttention对小目标的帮助更大,这符合它的设计初衷——通过局部上下文增强细节。
个人经验:什么时候该用,什么时候别用
推荐场景:
- 你的数据集里小目标占比高(比如自动驾驶、遥感图像)
- 模型已经足够轻量,想在不增加太多计算量的前提下提点
- 你发现加了SE或CA后mAP反而下降(这种情况我遇到过三次)
不推荐场景:
- 对推理速度要求极高(比如移动端实时检测),0.4ms的延迟在某些场景下不可接受
- 你的模型已经很大(比如YOLOv11x),再加注意力可能过拟合
- 数据集本身纹理简单(比如工业缺陷检测),局部上下文反而引入噪声
一个调试技巧:如果你发现加了CoTAttention后loss不下降,先检查dim是不是偶数。如果没问题,把sigmoid换成tanh试试,有时候梯度流会更顺畅。
最后说一句:注意力机制不是越多越好。我见过有人把SE、CBAM、CA、CoT全堆在一个模型里,结果mAP掉了2个点。少即是多,选一个最适合你数据集的,比堆砌一堆模块更有效。