学习资源站

RT-DETR改进策略【卷积层】HWD,引入`Haar小波变换`到下采样模块中,减少信息丢失_小波下采样-

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