RT-DETR改进策略【卷积层】| HWD,引入`Haar小波变换`到下采样模块中,减少信息丢失
一、本文介绍
本文记录的是
利用Haar小波下采样对RT-DETR网络进行改进的方法研究
。传统的卷积神经网络中常用的
最大池化
、
平均池化
和
步长为2的卷积
等操作进行下采样可能会导致信息丢失,为了解决信息丢失问题,
HWD
作者受无损信息变换方法的启发,引入
Haar小波变换
到下采样模块中,
旨在尽可能地保留图像信息,以便后续层能够提取更具判别性的特征,从而提高模型性能。
二、Haar小波下采样原理
Haar
小波下采样:一个简单但有效的语义分割下采样模块。
2.1、原理
HWD模
块由两个主要块组成:
无损特征编码块
和
特征表示学习块
。
-
无损特征编码块
:利用
Haar小波变换层有效地降低特征图的空间分辨率,同时保留所有信息。Haar小波变换是一种广泛认可的、紧凑的、二进的和正交的变换,在图像编码、边缘提取和二进制逻辑设计中有着广泛的应用。当对二维信号(如灰度图像)应用Haar小波变换时,会产生四个分量,每个分量的空间分辨率是原始信号的一半,而特征图的通道数则变为原来的四倍。这意味着Haar小波变换可以将部分空间维度的信息编码到通道维度中,而不会丢失任何信息。 -
特征表示学习块
:由标准的
1×1卷积层、批量归一化层和ReLU激活函数组成。该块用于调整特征图的通道数,使其与后续层对齐,并尽可能地过滤冗余信息,使后续层能够更有效地学习代表性特征。
2.2、优势
-
提高分割性能
:通过在三个不同模态的图像数据集上进行的广泛实验表明,
HWD模块能够有效提高分割性能。在Camvid数据集上,与七种最先进的分割架构相结合,使用HWD模块的模型在平均交并比(mIoU)上相比基线有1 - 2%的提升,特别是对于小尺度对象(如行人、自行车、围栏和标志符号等)的性能有显著改善。 - 减少信息不确定性 :利用结构相似性(SSIM)、峰值信噪比(PSNR)和提出的特征熵指数(FEI)评估下采样对特征图的有效性,结果表明HWD模块能够提高SSIM(7.78%)和PSNR(2.14 dB),并大幅降低信息不确定性。在所有21个模型中,HWD模块相比原始下采样方法,使特征不确定性降低了58.2%(FEI)和46.8%(FEI_B)。
-
通用性和易用性
:
HWD模块可以直接替换现有分割架构中的现有下采样方法(如最大池化、平均池化或步幅卷积),而不会引入额外的复杂性,并且能够显著提高分割性能。 -
在参数和计算量上的平衡
:与传统的下采样方法(如平均池化和步幅卷积)相比,
HWD模块在参数和浮点运算(FLOPs)上提供了一种平衡。虽然平均池化在参数和FLOPs方面表现更好,但HWD模块所需的参数少于步幅卷积的两倍,并且当通道数C大于一时,步幅卷积的计算开销超过HWD模块。 -
对浅层CNN的有效性
:在MOST数据集上的实验表明,当使用ResNet - 18和ResNet - 34作为特征提取的骨干网络时,HWD模块显著提高了分割性能,这表明浅层CNN对信息的需求更高,而
HWD模块能够满足这种需求。
HWD模块与其他下采样模块对比 :
保留信息能力
:传统的下采样方法(如最大池化、平均池化和步幅卷积等)会导致信息丢失,而
HWD模块
通过引入
Haar小波变换
,能够在降低特征图空间分辨率的同时尽可能保留信息。
论文: https://doi.org/10.1016/j.patcog.2023.109819
源码: https://github.com/apple1986/HWD
三、HWD的实现代码
HWD模块
的实现代码如下:
class HWD(nn.Module):
def __init__(self, in_ch, out_ch):
super(HWD, self).__init__()
from pytorch_wavelets import DWTForward
self.wt = DWTForward(J=1, mode='zero', wave='haar')
self.conv = Conv(in_ch * 4, out_ch, 1, 1)
def forward(self, x):
yL, yH = self.wt(x)
y_HL = yH[0][:, :, 0, ::]
y_LH = yH[0][:, :, 1, ::]
y_HH = yH[0][:, :, 2, ::]
x = torch.cat([yL, y_HL, y_LH, y_HH], dim=1)
x = self.conv(x)
return x
四、创新模块
4.1 改进点⭐
模块改进方法
:
1️⃣ 加入
HWD模块
。
HWD模块
添加后如下:
注意❗:在
5.2和5.3小节
中需要声明的模块名称为:
HWD
。
五、添加步骤
5.1 修改ultralytics/nn/modules/block.py
此处需要修改的文件是
ultralytics/nn/modules/block.py
block.py中定义了网络结构的通用模块
,我们想要加入新的模块就只需要将模块代码放到这个文件内即可。
将
HWD
模块代码添加到此文件下。
5.2 修改ultralytics/nn/modules/ init .py
此处需要修改的文件是
ultralytics/nn/modules/__init__.py
__init__.py
文件中定义了所有模块的初始化,我们只需要将
block.py
中的新的模块命添加到对应的函数即可。
HWD
在
block.py
中实现,所有要添加在
from .block import
:
from .block import (
C1,
C2,
...
HWD
)
5.3 修改ultralytics/nn/modules/tasks.py
在
tasks.py
文件中,需要在两处位置添加各模块类名称。
首先:在函数声明中引入
HWD
其次:在
parse_model函数
中注册
HWD
模块
还需在此函数下添加如下代码:
elif m in (HWD,):
args = [ch[f], ch[f]]
5.4 修改ultralytics/cfg/default.yaml
将
ultralytics/cfg/default.yaml
中的
amp
设置为
False
六、yaml模型文件
6.1 模型改进⭐
在代码配置完成后,配置模型的YAML文件。
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例,在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-HWD.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-HWD.yaml
文件下,修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
骨干网络
中的
下采样模块DWConv
替换成
HWD模块
。
# 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, HWD, [128]] # 2-P3/8
- [-1, 6, HGBlock, [96, 512, 3]] # stage 2
- [-1, 1, HWD, [512]] # 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, HWD, [1024]] # 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)
七、成功运行结果
分别打印网络模型可以看到
HWD模块
已经加入到模型中,并可以进行训练了。
yolov11m-HWD :
rtdetr-l-HWD summary: 685 layers, 38,101,571 parameters, 38,101,571 gradients, 115.5 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 65792 ultralytics.nn.AddModules.HWD.HWD [128, 128]
3 -1 6 839296 ultralytics.nn.modules.block.HGBlock [128, 96, 512, 3, 6]
4 -1 1 1049600 ultralytics.nn.AddModules.HWD.HWD [512, 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 4196352 ultralytics.nn.AddModules.HWD.HWD [1024, 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-HWD summary: 685 layers, 38,101,571 parameters, 38,101,571 gradients, 115.5 GFLOPs