最近在调一个缺陷检测的小模型输入分辨率不高目标又细又碎主干到 P3 那一层的特征图上看过去基本是一片糊。我一开始以为是数据增强和 anchor 匹配的问题把 mosaic 关了、把正样本阈值放宽、把 loss 权重来回挪了三遍mAP 也就上下浮动零点几个点。真正把指标顶上去的那一步是把主干里 C3 的聚合结果送进一个 CoTAttention 模块做一次上下文重标定再把改动接进 MMYOLO 的训练管线跑通。这篇就把整件事从头到尾拆一遍CoTAttention 在算什么、为什么它适合密集小目标、在 MMYOLO 里怎么注册怎么改 backbone、超参要跟着动哪些、以及我踩过的那些“mAP 先掉一个点再慢慢涨回来”的坑。内容偏实战前置知识只要求你会写 PyTorch 的 nn.Module、看得懂 MMYOLO 的 YAML 配置、跑过一次tools/train.py。如果你正好在做小目标、遮挡、弱纹理这几类检测任务又不想整块换主干换主干意味着重新调 neck 和 head时间成本太高那这种“在现有主干里插一个模块”的路子会省很多事。我会给完整代码、完整配置、参数量的估算过程以及一份排查表照着抄基本能在一两个小时内跑起来第一版。1. 主干为什么需要更强的上下文聚合1.1 三个真实场景暴露的短板先说清楚问题不然容易变成“为了加模块而加模块”。第一个场景是密集小目标。主干到 P3 输出的步长是 8一个 16×16 像素的目标在特征图上只剩 2×2。这时候单个位置的响应很容易被邻域的强纹理带偏卷积核的局部感受野又只覆盖到 3×3网络没有机会去“看一眼周围是什么”。结果就是一排焊点里漏掉中间那个或者把背景里的一条划痕认成缺陷。第二个场景是遮挡和截断。目标被挡住一半的时候可见部分和不可见部分在特征图上是分离的标准的 3×3 卷积没法建立这两块区域之间的关联。你只能靠堆层数、靠更深的网络去扩大感受野但那样又会带来梯度和效率的问题。第三个场景是弱纹理。工业件表面、医疗影像、遥感图里的很多目标颜色和灰度跟背景几乎一致区别只在“结构的连续性”上——比如一条边应该直着走断了就是缺陷。这种判断必须依赖上下文单点卷积无能为力。这三个问题的共同点都不是“分类头不够强”而是主干在早期阶段没有把空间上下文聚合好后面再补很费劲。1.2 CoTAttention 的核心思路静态上下文加动态注意力CoTAttention 的名字里有 Contextual Transformer它最讨巧的地方在于没有推翻卷积而是让卷积和自注意力各干各的活。具体拆成三条支路。第一条是静态上下文对输入做一次 k×k 卷积默认 3×3这一步的作用是把邻域信息先编码进来得到一个“局部上下文图”。第二条是动态上下文把上一步得到的上下文图和原始输入在通道维拼起来用两个 1×1 卷积压出一个k*k*C的输出再 reshape 成每个空间位置上的 k×k 个注意力权重。第三条是 value 支路1×1 卷积做通道变换。然后关键的一步是把动态上下文权重在 k×k 这个邻域维度上做 softmax再对 value 的邻域特征做加权求和。这样每个位置输出的是“融合了周围 k×k 个邻居、且权重由内容决定”的特征。最后再把静态上下文和动态上下文加起来作为整体输出的一部分。用一句人话概括卷积负责“看邻居”注意力负责“决定哪个邻居更重要”两者串起来就是 CoTAttention。整个过程是通道独立的不像标准自注意力那样需要 C×C 的注意力矩阵所以参数量和显存都控制得住。1.3 为什么选它而不是 SE、CBAM 或大核卷积这里我给一个自己在实际选型时用的对照表。这些模块我都往主干里插过结论是这个表模块增加的核心能力参数量C256 时显存增量我的实际体会SE通道重标定约 0.03M极小几乎不掉速但对空间上下文毫无帮助小目标场景收益常在 0.3 点以内CBAM通道 空间注意力约 0.07M小空间注意力图只有 1 通道粒度太粗密集场景容易误抑制大核 depthwise5/7扩大局部感受野约 0.02M小便宜好用但权重是静态的遇到内容差异大的邻域没法自适应CoTAttention内容自适应的邻域聚合约 0.95M中等小目标提升明显代价是显存和 FPS标准自注意力如 window attn全局建模通常 2M大在 P3 那种 80×80 的特征图上开销劝退除非做窗口划分选 CoT 的理由很直接我要的是“邻域内的内容自适应聚合”而不是“全局依赖”。检测任务里一个目标通常只跟它身边的几十个像素强相关把它拉通到全图反而引入噪声而且部署侧也不友好。注意CoT 的参数量集中在 1×1 的 query 生成卷积上输出维度是k*k*C通道数越大增长越快。这是后面要控制插入位置的根本原因不是拍脑袋决定的。2. 模块实现从张量形状开始抠2.1 四个张量的走向必须先想清楚写这类模块我习惯先把形状画在纸上不然一定会在 reshape 上翻车。CoTAttention 在 forward 里会经历四个关键张量static_ctx形状B,C,H,W由 k×k 卷积得到代表局部编码后的上下文。v形状B,C,H,W由 1×1 卷积得到作为待聚合的 value。attn形状B, k*k*C, H, W先 reshape 成B,C,k*k,H*W代表每个通道在每个位置的 k×k 个邻域权重。v_nb形状B,C,k*k,H*W是把 value 按 k×k 邻域展开之后的结果。注意attn和v_nb的维度顺序必须严格对齐都是(C, k*k, H*W)。只要这两者顺序一致后面就是一次普通的加权求和加一次沿邻域维的求和不会出错。我第一版就是attn用了(C, HW, kk)而v_nb用了(C, kk, HW)结果 loss 正常下降但指标死活不动——因为广播之后算出来的是完全无意义的东西而且不报错。2.2 一份能直接跑的完整实现下面这版我按“形状一目了然”的原则写用F.unfold显式拿到邻域 value。它跟原论文的动态上下文聚合在语义上等价好处是任何一步的形状你都能在注释里对照上调试时不会瞎猜。import torch import torch.nn as nn import torch.nn.functional as F class CoTAttention(nn.Module): Contextual Transformer Attention通道独立不改变分辨率。 def __init__(self, channels, kernel_size3, reduction4, norm_cfgNone): super().__init__() assert kernel_size % 2 1, kernel_size 必须是奇数否则 padding 对齐不了 self.channels channels self.kernel_size kernel_size self.pad kernel_size // 2 # 静态上下文k x k 卷积把邻域信息先编码进来 self.key_embed nn.Sequential( nn.Conv2d(channels, channels, kernel_size, paddingself.pad, biasFalse), nn.BatchNorm2d(channels), nn.SiLU(inplaceTrue), ) # value 支路1x1 通道变换 self.value_embed nn.Sequential( nn.Conv2d(channels, channels, 1, biasFalse), nn.BatchNorm2d(channels), ) # 动态上下文拼接后压出 k*k*C 个权重 mid max(channels // reduction, 8) self.query_embed nn.Sequential( nn.Conv2d(2 * channels, mid, 1, biasFalse), nn.BatchNorm2d(mid), nn.SiLU(inplaceTrue), nn.Conv2d(mid, kernel_size * kernel_size * channels, 1, biasTrue), ) def forward(self, x): b, c, h, w x.shape kk self.kernel_size * self.kernel_size hw h * w # 1) 静态上下文 B,C,H,W static_ctx self.key_embed(x) # 2) value B,C,H,W - 邻域展开 B,C,kk,HW v self.value_embed(x) v_nb F.unfold(v, self.kernel_size, paddingself.pad) # B, C*kk, HW v_nb v_nb.view(b, c, kk, hw) # 3) 动态上下文权重 B,kk*C,H,W - B,C,kk,HW attn self.query_embed(torch.cat([static_ctx, x], dim1)) attn attn.view(b, c, kk, hw) # 4) 动态上下文邻域求和代表这一片有什么 dynamic_ctx v_nb.sum(dim2).view(b, c, h, w) # 5) 邻域软注意力加权代表哪个邻居更该看 attn attn.softmax(dim2) out (attn * v_nb).sum(dim2).view(b, c, h, w) # 6) 三路融合 return out static_ctx dynamic_ctx几个设计选择需要说明。归一化我用了BatchNorm2d而不是原论文里的 LayerNorm原因是检测任务 batch size 通常不小8 卡 × 16BN 更稳而且 BN 在多数推理引擎里都有成熟算子导出时不会卡在 LayerNorm 上。如果你因为是极小 batch 才想用 LN记得把norm_cfg抽成参数别硬编码。激活我统一用 SiLU跟 MMYOLO 里 YOLOv5 系列的默认激活保持一致。混用 ReLU 和 SiLU 在浅层问题不大但会让 BN 的统计分布在不同 stage 之间差异变大调 lr 的时候更难伺候。2.3 怎么塞进 C3我试过的三种改法直接往 backbone 里裸插一个 CoT 是最省事的但效果一般因为它前面缺少一个通道压缩的过程输入通道大、计算量高。我的做法是把它挂在 C3 内部。MMYOLO 的 C3 结构是这样的cv1把输入压到 hidden 通道然后过m一串 Bottleneckcv2是另一条捷径最后 cat 之后cv3融合。最自然的插入点是m输出之后、cat 之前对聚合结果做一次上下文重标定。这样只加一个模块粒度可控而且 hidden 通道通常只有输出通道的一半参数量省一半。from mmcv.cnn import ConvModule from mmengine.model import BaseModule from mmyolo.models.layers import C3 from mmyolo.models.utils import make_divisible class C3_CoT(C3): 在 C3 的聚合分支末尾插入 CoTAttention。 def __init__(self, in_channels, out_channels, n1, shortcutTrue, expand_ratio0.5, norm_cfgdict(typeBN, momentum0.03, eps0.001), act_cfgdict(typeSiLU, inplaceTrue), init_cfgNone): super().__init__( in_channels, out_channels, nn, shortcutshortcut, expand_ratioexpand_ratio, norm_cfgnorm_cfg, act_cfgact_cfg, init_cfginit_cfg) hidden make_divisible(out_channels * expand_ratio, 8) self.cot CoTAttention(hidden) def forward(self, x): y1 self.m(self.cv1(x)) y1 self.cot(y1) return self.cv3(torch.cat((y1, self.cv2(x)), dim1))另外两种改法我也试过简单交代一下取舍。第一种是把 Bottleneck 中间的 3×3 卷积整个换成 CoT感受野够大但每个 Bottleneck 都加一个模块参数量和显存直接翻倍P3 上会 OOM放弃。第二种是整个 stage 后面单独挂一个 CoT 做后处理实现最简单但梯度回传路径变长训练前期非常不稳收敛慢也放弃了。2.4 把参数量和算力账算清楚这部分很多人跳过结果配置一跑起来发现显存超了。我按kernel_size3, reduction4手算一遍通道数记作 Ckey_embed3×3 卷积9C²加 BN 的2Cvalue_embed1×1 卷积C²加 BN 的2Cquery_embed第一层2C × C/4 0.5C²第二层C/4 × 9C 2.25C²合计2.75C²加两个 BN 约C/2 × 4总计约12.75C²量级的权重。用这个公式核对一下不同通道通道 C大概参数量典型出现位置64约 0.05Mbackbone stage2widen 0.5128约 0.21Mbackbone stage3widen 0.5256约 0.84Mbackbone stage4widen 0.5512约 3.3M大模型 stage4widen 1.0从表里能直接看出结论在 widen 0.5 的 s 模型上往 stage4 插一个模块只增加不到 1M 参数可以接受但同一套改动搬到 l 或 x 模型上stage4 的通道是 1024单模块就是 13M那就必须重新评估。这就是为什么别人分享的配置你不能无脑套到自己的模型尺度上。算力方面卷积部分可以用12.75C²HW粗估 MACs。以 stage4 为例C256、HW20×20400得到约 326M MACs换算成 GFLOPs 大概是 0.65。这个量级跟一个小号的 C3 差不多不算离谱。但真正的开销不在卷积在unfold产生的中间张量上下一节细说。3. MMYOLO 接入注册、改主干、改配置3.1 文件放哪怎么注册MMYOLO 基于注册器管理模块最规范的做法是用custom_imports引入自己的文件文件里用MODELS.register_module()装饰器注册。目录结构我一般这样组织cot_yolo/ ├── __init__.py ├── attention.py # CoTAttention └── backbone.py # C3_CoT YOLOv5CSPDarknetCoTbackbone.py里的注册写法from mmyolo.registry import MODELS from mmyolo.models.backbones import YOLOv5CSPDarknet from .attention import CoTAttention MODELS.register_module() class C3_CoT(C3): ... MODELS.register_module() class YOLOv5CSPDarknetCoT(YOLOv5CSPDarknet): ...然后在配置文件的顶层加custom_imports dict( imports[cot_yolo.backbone], allow_failed_importsFalse, )注意allow_failed_importsFalse一定要写。默认是 True导入失败的时候会被静默吞掉然后报一个“某某类不在注册表里”的错误你会以为是注册装饰器写错了实际上是路径不对。这一个参数我至少浪费过半小时。另外cot_yolo这个目录要么放在项目根目录下能被 Python 找到要么在运行前把路径塞进PYTHONPATH。我习惯直接放在根目录跟tools平级最省心。3.2 重写 build_stage而不是从零写 backbone这是整个接入过程里我认为最关键的一个选择。很多人一上来就自己写一个完整的 backbone 类把 stem、四个 stage、SPPF 全部重写一遍。我不建议这么干原因是 MMYOLO 的YOLOv5CSPDarknet里有一堆细节deepen_factor和widen_factor的取整逻辑make_divisible、每个 stage 的 stride 控制、输出通道记录、init_cfg的初始化策略。你手写的时候漏掉任何一个都会在 neck 那边报通道不匹配而且报错位置离真正的错误很远。我的做法是继承只重写build_stage
