RT-DETR改进策略【Neck】| ECCV-2024 RCM 矩形自校准模块 优化颈部网络
一、本文介绍
本文记录的是
利用
自校准模块RCM
优化RT-DETR的目标检测方法研究
。
RCM
通过矩形自校准函数可以将注意力区域校准得更接近前景对象，有效提高对前景对象的定位能力。
本文将其应用在颈部网络上，使模型能够捕获轴向全局上下文信息，并应用于金字塔上下文提取，使模型表现出更高的精度。
二、RCM 原理介绍
2.1 出发点
-
现有轻量级模型在特征表示能力上受限，难以对前景对象的边界进行建模和区分类别，导致边界分割不准确和分类错误。为了解决这些问题，设计了矩形
自校准模块（RCM）来提高前景对象的位置建模能力，并引入金字塔上下文来改善特征表示。
2.2 原理
2.2.1 矩形自校准注意力（RCA）
- 采用水平池化和垂直池化来捕获轴向全局上下文，生成两个轴向量。将这两个轴向量相加来建模一个矩形注意力区域。
- 设计形状自校准函数，通过大核条状卷积调整矩形注意力的形状，使其更接近前景特征。
2.2.2 特征融合
- 设计融合函数，将注意力特征与输入特征融合，使用 3 × 3 3×3 3 × 3 深度卷积进一步提取输入特征的局部细节，通过哈达玛积将校准后的注意力特征加权到细化后的输入特征上。
2.3 结构
-
由
矩形自校准注意力、批量归一化（BN）和多层感知机（MLP）组成。 - 矩形自校准注意力通过水平和垂直池化操作后，经过形状自校准函数校准，再进行特征融合。之后添加BN和MLP来细化特征，最后采用残差连接增强特征重用。
2.4 优势
-
位置建模和前景聚焦
- 能够使模型更专注于前景进行空间特征重建，通过形状自校准函数可以将注意力区域校准得更接近前景对象，有效提高对前景对象的定位能力。
-
上下文提取
- 在捕获轴向全局上下文用于金字塔上下文提取方面表现出色。通过水平和垂直池化以及后续的操作，可以更好地捕捉图像中的上下文信息。
-
性能优势
- 与现有的注意力机制相比，RCM通过其独特的设计，如形状自校准和特征融合等操作，可以取得更好的性能。例如在ADE20K数据集上的实验结果显示，使用RCM的模型在mIoU等指标上表现优异。
论文：h ttps://arxiv.org/pdf/2405.06228
源码： https://github.com/nizhenliang/CGRSeg
三、RCM的实现代码
RCM模块
的实现代码如下：
import torch
import torch.nn as nn
from timm.models.layers import DropPath, to_2tuple
from ultralytics.utils.torch_utils import fuse_conv_and_bn
class ConvMlp(nn.Module):
""" 使用 1x1 卷积保持空间维度的 MLP
"""
def __init__(
self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU,
norm_layer=None, bias=True, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
bias = to_2tuple(bias)
self.fc1 = nn.Conv2d(in_features, hidden_features, kernel_size=1, bias=bias[0])
self.norm = norm_layer(hidden_features) if norm_layer else nn.Identity()
self.act = act_layer()
self.drop = nn.Dropout(drop)
self.fc2 = nn.Conv2d(hidden_features, out_features, kernel_size=1, bias=bias[1])
def forward(self, x):
x = self.fc1(x)
x = self.norm(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
return x
#rectangular self-calibration attention (RCA)
class RCA(nn.Module):
def __init__(self, inp, kernel_size=1, ratio=1, band_kernel_size=11, dw_size=(1, 1), padding=(0, 0), stride=1,
square_kernel_size=2, relu=True):
super(RCA, self).__init__()
self.dwconv_hw = nn.Conv2d(inp, inp, square_kernel_size, padding=square_kernel_size // 2, groups=inp)
self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
self.pool_w = nn.AdaptiveAvgPool2d((1, None))
gc = inp // ratio
self.excite = nn.Sequential(
nn.Conv2d(inp, gc, kernel_size=(1, band_kernel_size), padding=(0, band_kernel_size // 2), groups=gc),
nn.BatchNorm2d(gc),
nn.ReLU(inplace=True),
nn.Conv2d(gc, inp, kernel_size=(band_kernel_size, 1), padding=(band_kernel_size // 2, 0), groups=gc),
nn.Sigmoid()
)
def sge(self, x):
# [N, D, C, 1]
x_h = self.pool_h(x)
x_w = self.pool_w(x)
x_gather = x_h + x_w # .repeat(1,1,1,x_w.shape[-1])
ge = self.excite(x_gather) # [N, 1, C, 1]
return ge
def forward(self, x):
loc = self.dwconv_hw(x)
att = self.sge(x)
out = att * loc
return out
#Rectangular Self-Calibration Module (RCM)
class RCM(nn.Module):
""" MetaNeXtBlock 块
参数:
dim (int): 输入通道数.
drop_path (float): 随机深度率。默认: 0.0
ls_init_value (float): 层级比例初始化值。默认: 1e-6.
"""
def __init__(
self,
dim,
token_mixer=RCA,
norm_layer=nn.BatchNorm2d,
mlp_layer=ConvMlp,
mlp_ratio=2,
act_layer=nn.GELU,
ls_init_value=1e-6,
drop_path=0.,
dw_size=11,
square_kernel_size=3,
ratio=1,
):
super().__init__()
self.token_mixer = token_mixer(dim, band_kernel_size=dw_size, square_kernel_size=square_kernel_size,
ratio=ratio)
self.norm = norm_layer(dim)
self.mlp = mlp_layer(dim, int(mlp_ratio * dim), act_layer=act_layer)
self.gamma = nn.Parameter(ls_init_value * torch.ones(dim)) if ls_init_value else None
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
def forward(self, x):
shortcut = x
x = self.token_mixer(x)
x = self.norm(x)
x = self.mlp(x)
if self.gamma is not None:
x = x.mul(self.gamma.reshape(1, -1, 1, 1))
x = self.drop_path(x) + shortcut
return x
四、创新模块
4.1 改进点⭐
模块改进方法
：加入
RCM模块
（
第五节讲解添加步骤
）。
RCM模块
添加后如下：
五、添加步骤
5.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
RCM.py
，将
第三节
中的代码粘贴到此处
5.2 修改二
在
AddModules
文件夹下新建
__init__.py
（已有则不用新建），在文件内导入模块：
from .RCM import *
5.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中，需要在两处位置添加各模块类名称。
首先：导入模块
其次：在
parse_model函数
中添加如下代码：
elif m in {RCM}:
args = [ch[f]]
六、yaml模型文件
6.1 模型改进版本
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-RCM.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-RCM.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是在
颈部网络
中添加
RCM模块
。
# 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-P3/16
- [-1, 6, HGBlock, [192, 1024, 5, True, False]] # cm, c2, k, light, shortcut
- [-1, 6, HGBlock, [192, 1024, 5, True, True]]
- [-1, 6, HGBlock, [192, 1024, 5, True, True]] # stage 3
- [-1, 1, DWConv, [1024, 3, 2, 1, False]] # 8-P4/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, RCM, []] # 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, RCM, []] # X3 (21), fpn_blocks.1
- [-1, 1, Conv, [256, 3, 2]] # 22, downsample_convs.0
- [[-1, 17], 1, Concat, [1]] # cat Y4
- [-1, 3, RCM, []] # F4 (24), pan_blocks.0
- [-1, 1, Conv, [256, 3, 2]] # 25, downsample_convs.1
- [[-1, 12], 1, Concat, [1]] # cat Y5
- [-1, 3, RCM, []] # F5 (27), pan_blocks.1
- [[21, 24, 27], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5)
七、成功运行结果
打印网络模型可以看到
RCM
已经加入到模型中，并可以进行训练了。
rtdetr-l-RCM ：
rtdetr-l-RCM summary: 768 layers, 38,161,603 parameters, 38,161,603 gradients, 131.1 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 1695360 ultralytics.nn.modules.block.HGBlock [512, 192, 1024, 5, 6, True, False]
6 -1 6 2055808 ultralytics.nn.modules.block.HGBlock [1024, 192, 1024, 5, 6, True, True]
7 -1 6 2055808 ultralytics.nn.modules.block.HGBlock [1024, 192, 1024, 5, 6, True, True]
8 -1 1 11264 ultralytics.nn.modules.conv.DWConv [1024, 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 262656 ultralytics.nn.modules.conv.Conv [1024, 256, 1, 1, None, 1, 1, False]
15 [-2, -1] 1 0 ultralytics.nn.modules.conv.Concat [1]
16 -1 3 3210240 ultralytics.nn.AddModules.RCM.RCM [512]
17 -1 1 131584 ultralytics.nn.modules.conv.Conv [512, 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 3210240 ultralytics.nn.AddModules.RCM.RCM [512]
22 -1 1 1180160 ultralytics.nn.modules.conv.Conv [512, 256, 3, 2]
23 [-1, 17] 1 0 ultralytics.nn.modules.conv.Concat [1]
24 -1 3 3210240 ultralytics.nn.AddModules.RCM.RCM [512]
25 -1 1 1180160 ultralytics.nn.modules.conv.Conv [512, 256, 3, 2]
26 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]
27 -1 3 3210240 ultralytics.nn.AddModules.RCM.RCM [512]
28 [21, 24, 27] 1 7500515 ultralytics.nn.modules.head.RTDETRDecoder [1, [512, 512, 512]]
rtdetr-l-RCM summary: 768 layers, 38,161,603 parameters, 38,161,603 gradients, 131.1 GFLOPs