RT-DETR改进策略【卷积层】| 2024最新轻量级自适应提取模块 LAE 即插即用 保留局部信息和全局信息
一、本文介绍
本文记录的是
利用
轻量级自适应提取模块(LAE)
模块优化
RT-DETR
的目标检测网络模型
。
LAE (Lightweight Adaptive Extraction)
在
减少参数和计算成本
的同时，能够
提取更丰富语义信息的特征
，克服了传统卷积方法
难以捕捉全局信息
的问题，并能更好地提取
ROI
特征。本文将其应用
RT-DETR
中，改进主要模块，更好地突出重要特征，从而提升模型对物体检测的能力。
二、轻量级自适应提取（LAE）模块介绍
2.1 设计出发点
- 减少参数和计算成本 ：在多尺度特征提取中，相较于传统卷积方法，旨在降低参数数量和计算成本，同时提取具有更丰富语义信息的特征。
- 保留局部信息 ：传统卷积过程中边缘和角落的像素信息易丢失，而这些局部信息在医学成像中对表示ROI目标和周围区域的隐含信息尤为重要。同时，特征图中相邻像素间存在信息差异，包含目标的像素信息熵高于相邻像素，所以希望在采样过程中保留信息含量高的像素。
- 解决卷积局限性 ：考虑到卷积操作的局部性难以捕捉全局信息且计算复杂，参考了Focus的概念，但因Focus的切片操作计算成本高不符合轻量化处理目标，所以在此基础上进行改进。
2.2 原理
- 通过设计两个共享参数的并行分支，采用组卷积的概念高效地将输入映射到输出维度，利用N组组卷积使参数数量相比传统卷积减少到1/N。
- 每个LAE单元实现四倍下采样（高度和宽度都缩小为原来的1/2）。在采样过程中，为缓解边缘信息丢失问题，先将特征图的高度和宽度信息保存到通道中，使特征图维度从四维（batch, channel, height, weight）变为五维（batch, channel, height, weight, n），其中‘n’表示采样因子。
- 自适应提取路径通过平均池化和卷积交换信息，对特征图按四个相邻像素（如左上角四个像素）重新组合进行下采样，其权重通过softmax表示，维度也变为五维。在‘n’维度上，自适应权重与另一分支相结合，这种方式可理解为在从高分辨率到低分辨率信息的转换过程中，在通道层面隐式包含了全局信息。
2.3 结构
- 由两个并行分支构成，包括自适应提取路径和轻量化提取路径。
- 自适应提取路径 ：通过平均池化和卷积交换信息，实现特征图按四个相邻像素重新组合下采样，并确定自适应权重。
- 轻量化提取路径 ：利用1×1卷积、平均池化、重排操作、Softmax乘法等操作，高效地将输入映射到输出维度。
2.4 优势
- 参数少且易使用 ：模块使用时无参数，便于即插即用式的迁移。
- 保留信息丰富 ：能有效保留局部信息和在采样过程中隐式包含全局信息，从而提取出具有更丰富语义信息的特征。
论文： https://arxiv.org/pdf/2408.14087
源码： https://github.com/VincentYuuuuuu/LSM-YOLO
三、LAE的实现代码
LAE
的实现代码如下：
import torch
import torch.nn as nn
from einops import rearrange
def autopad(k, p=None, d=1): # kernel, padding, dilation
"""Pad to 'same' shape outputs."""
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):
"""Initialize Conv layer with given arguments including 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):
"""Apply convolution, batch normalization and activation to input tensor."""
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
"""Perform transposed convolution of 2D data."""
return self.act(self.conv(x))
class LAE(nn.Module):
# Light-weight Adaptive Extraction
def __init__(self, ch, group=8) -> None:
super().__init__()
self.softmax = nn.Softmax(dim=-1)
self.attention = nn.Sequential(
nn.AvgPool2d(kernel_size=3, stride=1, padding=1),
Conv(ch, ch, k=1)
)
self.ds_conv = Conv(ch, ch * 4, k=3, s=2, g=(ch // group))
def forward(self, x):
# bs, ch, 2*h, 2*w => bs, ch, h, w, 4
att = rearrange(self.attention(x), 'bs ch (s1 h) (s2 w) -> bs ch h w (s1 s2)', s1=2, s2=2)
att = self.softmax(att)
# bs, 4 * ch, h, w => bs, ch, h, w, 4
x = rearrange(self.ds_conv(x), 'bs (s ch) h w -> bs ch h w s', s=4)
x = torch.sum(x * att, dim=-1)
return x
四、创新模块
4.1 改进点⭐
模块改进方法
：加入
LAE模块
（
第五节讲解添加步骤
）。
LAE模块
添加后如下：
五、添加步骤
5.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
LAE.py
，将
第三节
中的代码粘贴到此处
5.2 修改二
在
AddModules
文件夹下新建
__init__.py
（已有则不用新建），在文件内导入模块：
from .LAE import *
5.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中，需要在两处位置添加各模块类名称。
首先：导入模块
其次：在
parse_model函数
中添加如下代码：
elif m in {LAE}:
c2 = ch[f]
args = [c2, *args]
六、yaml模型文件
6.1 模型改进版本
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-LAE.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-LAE.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
骨干网络
中的 下采样模块替换成
LAE模块
。
# 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, LAE, []] # 2-P3/8
- [-1, 6, HGBlock, [96, 512, 3]] # stage 2
- [-1, 1, LAE, []] # 4-P4/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, LAE, []] # 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)
七、成功运行结果
打印网络模型可以看到
LAE
已经加入到模型中，并可以进行训练了。
rtdetr-l-LAE ：
rtdetr-l-LAE summary: 700 layers, 34,612,803 parameters, 34,612,803 gradients, 116.9 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 54528 ultralytics.nn.AddModules.LAE.LAE [128]
3 -1 6 839296 ultralytics.nn.modules.block.HGBlock [128, 96, 512, 3, 6]
4 -1 1 414720 ultralytics.nn.AddModules.LAE.LAE [512]
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 1353728 ultralytics.nn.AddModules.LAE.LAE [1024]
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 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-LAE summary: 700 layers, 34,612,803 parameters, 34,612,803 gradients, 116.9 GFLOPs