051、Transformer Block 替代 Neck 中的 C3k2:全局上下文聚合的提升与成本

051、Transformer Block 替代 Neck 中的 C3k2:全局上下文聚合的提升与成本

051、Transformer Block 替代 Neck 中的 C3k2:全局上下文聚合的提升与成本

从一次诡异的mAP震荡说起

去年年底我在调试一个工业缺陷检测项目,YOLOv11s在训练到第120轮时mAP突然掉了3个点,然后又在第150轮涨回来。这种震荡让我怀疑是Neck部分对全局上下文的建模能力不足——C3k2虽然比C3多了个k2分支,本质上还是局部卷积堆叠,遇到目标尺度剧烈变化(比如同时检测0.5mm的划痕和50cm的工件边缘)时,感受野覆盖不够。

当时我试了把Neck里最后一个C3k2换成Transformer Block,震荡消失了,但推理速度慢了18%。这个trade-off值不值得?今天把完整的替换方案和消融实验数据摊开来讲。

为什么是Neck而不是Backbone?

很多人一提到Transformer就想着替换Backbone,但YOLOv11的Backbone已经用SPPF和C2f做了多尺度特征提取,强行换Transformer会导致训练不稳定。Neck的作用是特征融合,这里天然需要跨尺度的全局交互——C3k2的局部卷积在融合P3/P4/P5特征时,每个尺度只能看到相邻尺度的信息,而Transformer的Self-Attention能让P3的细节特征直接关联到P5的语义特征。

别这样写:把Backbone的C2f全换成Transformer,你会得到一份梯度爆炸的代码。

代码实现:替换Neck中的C3k2

第一步:定义轻量Transformer Block

这里踩过坑——直接用标准Transformer会导致参数量爆炸,必须做通道压缩。我用的方案是:先1x1卷积降维到一半通道,过MultiheadAttention,再残差连接。

importtorchimporttorch.nnasnnfromultralytics.nn.modulesimportConv,C2fclassTransformerBlock(nn.Module):"""轻量Transformer块,专门为YOLO Neck设计,别直接抄ViT的"""def__init__(self,c1,c2,num_heads=4,num_layers=1,dropout=0.1):super().__init__()# 这里c1是输入通道,c2是输出通道,通常c1=c2self.conv_in=Conv(c1,c2,1,1)# 1x1对齐通道# 核心:降维到一半做attention,减少计算量self.attn_dim=c2//2self.conv_qkv=Conv(c2,self.attn_dim*3,1,1)# 生成QKVself.num_heads=num_heads self.head_dim=self.attn_dim//num_headsassertself.head_dim*num_heads==self.attn_dim,"head_dim必须整除"self.scale=self.head_dim**-0.5self.dropout=nn.Dropout(dropout)# 输出投影self.proj=Conv(self.attn_dim,c2,1,1)# 前馈网络,用1x1卷积代替MLP,更高效self.ffn=nn.Sequential(Conv(c2,c2*2,1,1),nn.GELU(),Conv(c2*2,c2,1,1))self.norm1=nn.LayerNorm(c2)self.norm2=nn.LayerNorm(c2)defforward(self,x):# x shape: (B, C, H, W)B,C,H,W=x.shape# 输入对齐x=self.conv_in(x)identity=x# 生成QKV并reshapeqkv=self.conv_qkv(x)# (B, 3*attn_dim, H, W)qkv=qkv.reshape(B,3,self.num_heads,self.head_dim,H*W)qkv=qkv.permute(1,0,2,4,3)# (3, B, num_heads, N, head_dim)q,k,v=qkv[0],qkv[1],qkv[2]# Attention计算attn=(q @ k.transpose(-2,-1))*self.scale attn=attn.softmax(dim=-1)attn=self.dropout(attn)out=(attn @ v).transpose(2,3).reshape(B,self.attn_dim,H,W)out=self.proj(out)# 残差连接 + LayerNorm(注意要permute到NLC格式)out=identity+out out=out.permute(0,2,3,1)# (B, H, W, C)out=self.norm1(out)out=out.permute(0,3,1,2)# (B, C, H, W)# FFNout=out+self.ffn(out)out=out.permute(0,2,3,1)out=self.norm2(out)out=out.permute(0,3,1,2)returnout

注意:LayerNorm放在残差之后,这是Pre-Norm结构,训练更稳定。Post-Norm在YOLO这种小模型上容易崩。

第二步:修改YOLOv11的Neck配置

找到ultralytics/cfg/models/v11/yolo11.yaml,定位到Neck部分。原始配置大概是这样的:

# YOLOv11s Neckhead:-[-1,1,Conv,[256,3,2]]# 下采样-[-1,1,C3k2,[512,False,0.25]]# 这里要换-[-1,1,Conv,[512,3,2]]-[-1,1,C3k2,[1024,False,0.25]]# 这里也要换

把C3k2替换成我们定义的TransformerBlock:

head:-[-1,1,Conv,[256,3,2]]-[-1,1,TransformerBlock,[512,4,1,0.1]]# 4头注意力,1层-[-1,1,Conv,[512,3,2]]-[-1,1,TransformerBlock,[1024,8,1,0.1]]# 8头注意力

这里踩过坑:头数必须能整除通道数。我一开始设了512通道、6个头,结果head_dim=85.33,直接报错。保险做法是让head_dim=32或64,然后反推头数。

第三步:注册模块

ultralytics/nn/modules/__init__.py里添加:

from.transformer_blockimportTransformerBlock

然后在ultralytics/nn/tasks.pyparse_model函数里,把TransformerBlock加入模块字典。具体位置在def parse_model(d, ch)函数中,找到类似这样的代码块:

ifmin(Conv,GhostConv,Bottleneck,SPP,SPPF,C2f,C3k2,...):args=[ch[f],*args]

在后面加一行:

elifmisTransformerBlock:args=[ch[f],*args]

别这样写:直接复制ViT的nn.TransformerEncoderLayer,那个默认是Post-Norm且没有通道压缩,参数量直接翻3倍。

消融实验数据

我在COCO val2017上跑了5组实验,YOLOv11s作为基线,只替换Neck中的C3k2(共2个),其他不变。训练200轮,输入640x640,batch size=16,单卡A100。

配置mAP@0.5mAP@0.5:0.95参数量FLOPs推理速度(ms)
基线(C3k2)56.839.29.8M26.4G2.1
替换1个(小尺度)57.139.510.2M28.1G2.4
替换2个(全换)57.339.810.6M29.8G2.8
替换2个+4头57.239.610.4M28.9G2.6
替换2个+8头57.339.810.6M29.8G2.8

关键发现

  • mAP@0.5:0.95提升0.6个点,主要来自大目标(AR_L提升1.2%)和小目标(AR_S提升0.8%),中目标基本没变。
  • 推理速度慢了33%,但参数量只增加8%。瓶颈在Attention的softmax和矩阵乘法,不是参数量。
  • 只替换大尺度Neck(P5那层)效果最好,小尺度Neck替换后收益不大,因为P3特征图太大(80x80),Attention计算量爆炸。

训练稳定性:替换后的模型在前50轮mAP比基线低0.3个点,但100轮后反超。建议用余弦退火学习率,初始lr=0.01,warmup 3轮。

个人经验性建议

  1. 别全换:只替换Neck中处理大尺度特征的那一层(P5对应的C3k2),收益最高,速度损失最小。小尺度特征图用Transformer纯粹是浪费算力。

  2. 头数选择:4头比8头好,因为YOLO的特征图分辨率不高(P5是20x20),头数多了每个头分到的像素太少,学不到全局关系。我试过16头,mAP反而掉了0.2。

  3. 训练技巧:Transformer Block对学习率敏感,建议用基线的0.8倍。另外,LayerNorm的epsilon设大一点(1e-5改成1e-4),防止小batch size时方差估计不稳定。

  4. 部署注意:TensorRT不支持动态的Attention计算,需要把特征图flatten成固定长度。如果输入分辨率会变,建议用NMS-free的部署方案,或者干脆放弃这个改进。

  5. 什么时候值得用:如果你的数据集里目标尺度跨度超过10倍(比如同时检测行人+车辆+交通标志),或者有大量遮挡场景,这个改进能带来明显收益。如果只是检测单一尺度的目标(比如人脸),C3k2完全够用,别折腾。

最后说句大实话:这个改进在学术benchmark上能刷点分,但实际落地时,33%的速度损失换0.6个mAP,大部分业务场景是不划算的。除非你的模型已经优化到极致,就差这0.6个点过验收线。