RT-DETR改进策略【Conv和Transformer】| 上下文转换器CoT 结合静态和动态上下文信息的注意力机制 (含二次创新)
一、本文介绍
本文记录的是
利用
Contextual Transformer (CoT)
模块优化
RT-DETR
的目标检测网络模型
。
CoT
将静态上下文和自注意力学习动态上下文统一在一个架构中，
有效地提升了在 2D 特征图上进行视觉表示学习时自注意力学习的能力
，本文将深入研究
CoT
的原理，并将其应用到
RT-DETR
中，
进而更有效的增强模型对视觉信息的学习和表示能力。
二、ParNet 介绍
Contextual Transformer Networks for Visual Recognition
2.1 出发点
- 现有的大多数Transformer设计在2D特征图上直接使用自注意力来获得注意力矩阵时，仅基于孤立的查询和键对，未充分利用相邻键之间的丰富上下文信息。因此，希望通过利用输入键之间的上下文信息来增强Transformer式架构，以提高视觉表示学习能力。
2.2 原理
- 挖掘静态上下文 ：首先通过一个 3 × 3 3×3 3 × 3 卷积对输入键进行上下文编码，得到输入的静态上下文表示。此操作考虑了每个 3 × 3 3×3 3 × 3 网格内相邻键之间的关系，反映了局部邻居之间的静态上下文。
- 学习动态注意力矩阵 ：将编码后的键与输入查询连接起来，通过两个连续的 1 × 1 1×1 1 × 1 卷积来学习动态多头注意力矩阵。这个过程在静态上下文的指导下，自然地利用了每个查询和所有键之间的相互关系进行自注意力学习。
- 获得动态上下文表示 ：将学习到的注意力矩阵与输入值相乘，得到输入的动态上下文表示，描绘了动态上下文。
- 融合输出 ：最后将静态和动态上下文表示融合作为输出。
2.3 结构
-
如图所示，假设输入的2D特征图为 X X X ，键 K = X K = X K = X ，查询 Q = X Q = X Q = X ，值 V = X W v V = XW_{v} V = X W v 。首先使用 k × k k×k k × k 组卷积对 k × k k×k k × k 网格内的相邻键进行上下文编码得到 K 1 K^{1} K 1 ，作为静态上下文表示。
-
然后基于 [ K 1 , Q ] [K^{1}, Q] [ K 1 , Q ] 通过两个连续的 1 × 1 1×1 1 × 1 卷积（ W θ W_{\theta} W θ 带有ReLU激活函数， W δ W_{\delta} W δ 无激活函数）得到注意力矩阵 A A A 。
-
接着根据注意力矩阵 A A A 计算得到动态上下文表示 K 2 = V ⊛ A K^{2}=V \circledast A K 2 = V ⊛ A 。
-
最终输出 Y Y Y 是静态上下文 K 1 K^{1} K 1 和动态上下文 K 2 K^{2} K 2 通过注意力机制融合的结果。
-
优势
- 统一架构 ：将键之间的上下文挖掘和2D特征图上的自注意力学习统一在一个架构中，避免了为上下文挖掘引入额外的分支，同时具有良好的参数预算。
- 有效利用上下文 ：充分利用了邻居键之间的上下文信息，增强了自注意力学习，提高了输出聚合特征图的表示能力。
-
可替代卷积
：可以作为现有ResNet架构中标准卷积的替代，直接替换ResNet结构中的
3
×
3
3×3
3
×
3
卷积，得到名为
Contextual Transformer Networks (CoTNet)的新架构，且在不增加参数和FLOP预算的情况下提升性能。例如，在ImageNet图像识别任务中，CoTNet相对于ResNeSt (101层)，top - 1错误率绝对降低了 0.9 % 0.9\% 0.9% ；在COCO的目标检测和实例分割任务中，分别绝对提高了ResNeSt的 1.5 % 1.5\% 1.5% 和 0.7 % 0.7\% 0.7% 的mAP。
论文： https://arxiv.org/pdf/2107.12292
源码： https://github.com/JDAI-CV/CoTNet
三、CoT的实现代码
CoT模块
的实现代码如下：
import torch
import torch.nn as nn
import torch.nn.functional as F
from ultralytics.nn.modules.conv import LightConv
from ultralytics.utils.torch_utils import fuse_conv_and_bn
class CoTAttention(nn.Module):
def __init__(self, dim=512, kernel_size=3):
super().__init__()
self.dim = dim
self.kernel_size = kernel_size
self.key_embed = nn.Sequential(
nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=4, bias=False),
nn.BatchNorm2d(dim),
nn.SiLU()
)
self.value_embed = nn.Sequential(
nn.Conv2d(dim, dim, 1, bias=False),
nn.BatchNorm2d(dim)
)
factor = 4
self.attention_embed = nn.Sequential(
nn.Conv2d(2 * dim, 2 * dim // factor, 1, bias=False),
nn.BatchNorm2d(2 * dim // factor),
nn.SiLU(),
nn.Conv2d(2 * dim // factor, kernel_size * kernel_size * dim, 1)
)
def forward(self, x):
bs, c, h, w = x.shape
k1 = self.key_embed(x) # bs,c,h,w
v = self.value_embed(x).view(bs, c, -1) # bs,c,h,w
y = torch.cat([k1, x], dim=1) # bs,2c,h,w
att = self.attention_embed(y) # bs,c*k*k,h,w
att = att.reshape(bs, c, self.kernel_size * self.kernel_size, h, w)
att = att.mean(2, keepdim=False).view(bs, c, -1) # bs,c,h*w
k2 = F.softmax(att, dim=-1) * v
k2 = k2.view(bs, c, h, w)
return k1 + k2
def autopad(k, p=None, d=1):
"""
Pads kernel to 'same' output shape, adjusting for optional dilation; returns padding size.
`k`: kernel, `p`: padding, `d`: dilation.
"""
if d > 1:
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
class Conv(nn.Module):
# Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
default_act = nn.SiLU() # default activation
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
"""Initializes a standard convolution layer with optional batch normalization and activation."""
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self, x):
"""Applies a convolution followed by batch normalization and an activation function to the input tensor `x`."""
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
"""Applies a fused convolution and activation function to the input tensor `x`."""
return self.act(self.conv(x))
class HGBlock_CoT(nn.Module):
"""
HG_Block of PPHGNetV2 with 2 convolutions and LightConv.
https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
"""
def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
"""Initializes a CSP Bottleneck with 1 convolution using specified input and output channels."""
super().__init__()
block = LightConv if lightconv else Conv
self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act) # squeeze conv
self.ec = Conv(c2 // 2, c2, 1, 1, act=act) # excitation conv
self.add = shortcut and c1 == c2
self.cv = CoTAttention(c2)
def forward(self, x):
"""Forward pass of a PPHGNetV2 backbone layer."""
y = [x]
y.extend(m(y[-1]) for m in self.m)
y = self.cv(self.ec(self.sc(torch.cat(y, 1))))
return y + x if self.add else y
四、创新模块
4.1 改进点⭐
模块改进方法
：加入
CoTAttention模块
（
第五节讲解添加步骤
）。
CoTAttention模块
添加后如下：
基于
CoTAttention模块
的
HGBlock
（
第五节讲解添加步骤
）。
此时是对
RT-DETR
中的
HGBlock模块
进行改进，并将
CoTAttention
在加入到
HGBlock
模块中。
改进代码如下：
首先，将
CoTAttention
添加到
HGBlock
模块中
class HGBlock_CoT(nn.Module):
"""
HG_Block of PPHGNetV2 with 2 convolutions and LightConv.
https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
"""
def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
"""Initializes a CSP Bottleneck with 1 convolution using specified input and output channels."""
super().__init__()
block = LightConv if lightconv else Conv
self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act) # squeeze conv
self.ec = Conv(c2 // 2, c2, 1, 1, act=act) # excitation conv
self.add = shortcut and c1 == c2
self.cv = CoTAttention(c2)
def forward(self, x):
"""Forward pass of a PPHGNetV2 backbone layer."""
y = [x]
y.extend(m(y[-1]) for m in self.m)
y = self.cv(self.ec(self.sc(torch.cat(y, 1))))
return y + x if self.add else y
注意❗：在
第五小节
中需要声明的模块名称为：
HGBlock_CoT
。
五、添加步骤
5.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
CoT.py
，将
第三节
中的代码粘贴到此处
5.2 修改二
在
AddModules
文件夹下新建
__init__.py
（已有则不用新建），在文件内导入模块：
from .CoT import *
5.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中，需要在两处位置添加各模块类名称。
首先：导入模块
其次：在
parse_model函数
中注册
HGBlock_CoT
模块
最后，在此函数下添加如下代码：
elif m is CoTAttention:
c1, c2 = ch[f], args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels) * width, 8)
args = [c1, *args[1:]]
六、yaml模型文件
6.1 模型改进版本⭐
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-HGBlock_CoT.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-HGBlock_CoT.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
网络
中的所有
HGBlock模块
替换成
HGBlock_CoT模块
。
# Ultralytics YOLO 🚀, AGPL-3.0 license
# RT-DETR-l object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/rtdetr
# Parameters
nc: 1 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n'
# [depth, width, max_channels]
l: [1.00, 1.00, 1024]
backbone:
# [from, repeats, module, args]
- [-1, 1, HGStem, [32, 48]] # 0-P2/4
- [-1, 6, HGBlock, [48, 128, 3]] # stage 1
- [-1, 1, DWConv, [128, 3, 2, 1, False]] # 2-P3/8
- [-1, 6, HGBlock, [96, 512, 3]] # stage 2
- [-1, 1, DWConv, [512, 3, 2, 1, False]] # 4-P4/16
- [-1, 6, HGBlock_CoT, [192, 512, 5, True, False]] # cm, c2, k, light, shortcut
- [-1, 6, HGBlock_CoT, [192, 512, 5, True, True]]
- [-1, 6, HGBlock_CoT, [192, 512, 5, True, True]] # stage 3
- [-1, 1, DWConv, [1024, 3, 2, 1, False]] # 8-P5/32
- [-1, 6, HGBlock, [384, 2048, 5, True, False]] # stage 4
head:
- [-1, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 10 input_proj.2
- [-1, 1, AIFI, [1024, 8]]
- [-1, 1, Conv, [256, 1, 1]] # 12, Y5, lateral_convs.0
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [7, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 14 input_proj.1
- [[-2, -1], 1, Concat, [1]]
- [-1, 3, RepC3, [256]] # 16, fpn_blocks.0
- [-1, 1, Conv, [256, 1, 1]] # 17, Y4, lateral_convs.1
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [3, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 19 input_proj.0
- [[-2, -1], 1, Concat, [1]] # cat backbone P4
- [-1, 3, RepC3, [256]] # X3 (21), fpn_blocks.1
- [-1, 1, Conv, [256, 3, 2]] # 22, downsample_convs.0
- [[-1, 17], 1, Concat, [1]] # cat Y4
- [-1, 3, RepC3, [256]] # F4 (24), pan_blocks.0
- [-1, 1, Conv, [256, 3, 2]] # 25, downsample_convs.1
- [[-1, 12], 1, Concat, [1]] # cat Y5
- [-1, 3, RepC3, [256]] # F5 (27), pan_blocks.1
- [[21, 24, 27], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5)
七、成功运行结果
打印网络模型可以看到
HGBlock_CoT
已经加入到模型中，并可以进行训练了。
rtdetr-l-HGBlock_CoT ：
rtdetr-l-HGBlock_CoT summary: 721 layers, 36,396,739 parameters, 36,396,739 gradients, 119.4 GFLOPs
from n params module arguments
0 -1 1 25248 ultralytics.nn.modules.block.HGStem [3, 32, 48]
1 -1 6 155072 ultralytics.nn.modules.block.HGBlock [48, 48, 128, 3, 6]
2 -1 1 1408 ultralytics.nn.modules.conv.DWConv [128, 128, 3, 2, 1, False]
3 -1 6 839296 ultralytics.nn.modules.block.HGBlock [128, 96, 512, 3, 6]
4 -1 1 5632 ultralytics.nn.modules.conv.DWConv [512, 512, 3, 2, 1, False]
5 -1 6 3175552 ultralytics.nn.AddModules.CoT.HGBlock_CoT [512, 192, 512, 5, 6, True, False]
6 -1 6 3175552 ultralytics.nn.AddModules.CoT.HGBlock_CoT [512, 192, 512, 5, 6, True, True]
7 -1 6 3175552 ultralytics.nn.AddModules.CoT.HGBlock_CoT [512, 192, 512, 5, 6, True, True]
8 -1 1 11264 ultralytics.nn.modules.conv.DWConv [512, 1024, 3, 2, 1, False]
9 -1 6 6708480 ultralytics.nn.modules.block.HGBlock [1024, 384, 2048, 5, 6, True, False]
10 -1 1 524800 ultralytics.nn.modules.conv.Conv [2048, 256, 1, 1, None, 1, 1, False]
11 -1 1 789760 ultralytics.nn.modules.transformer.AIFI [256, 1024, 8]
12 -1 1 66048 ultralytics.nn.modules.conv.Conv [256, 256, 1, 1]
13 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
14 7 1 131584 ultralytics.nn.modules.conv.Conv [512, 256, 1, 1, None, 1, 1, False]
15 [-2, -1] 1 0 ultralytics.nn.modules.conv.Concat [1]
16 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
17 -1 1 66048 ultralytics.nn.modules.conv.Conv [256, 256, 1, 1]
18 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
19 3 1 131584 ultralytics.nn.modules.conv.Conv [512, 256, 1, 1, None, 1, 1, False]
20 [-2, -1] 1 0 ultralytics.nn.modules.conv.Concat [1]
21 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
22 -1 1 590336 ultralytics.nn.modules.conv.Conv [256, 256, 3, 2]
23 [-1, 17] 1 0 ultralytics.nn.modules.conv.Concat [1]
24 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
25 -1 1 590336 ultralytics.nn.modules.conv.Conv [256, 256, 3, 2]
26 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]
27 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
28 [21, 24, 27] 1 7303907 ultralytics.nn.modules.head.RTDETRDecoder [1, [256, 256, 256]]
rtdetr-l-HGBlock_CoT summary: 721 layers, 36,396,739 parameters, 36,396,739 gradients, 119.4 GFLOPs