一、本文介绍
本文记录的是利用Haar小波下采样对YOLOv9网络进行改进的方法研究。传统的卷积神经网络中常用的最大池化、平均池化和步长为2的卷积等操作进行下采样可能会导致信息丢失，为了解决信息丢失问题，HWD作者受无损信息变换方法的启发，引入Haar小波变换到下采样模块中，旨在尽可能地保留图像信息，以便后续层能够提取更具判别性的特征，从而提高分割性能。
文章目录
- 一、本文介绍
- 二、Haar小波下采样原理
- 2.1、原理
- 2.2、优势
- 三、HWD的实现代码
- 四、添加步骤
- 4.1 修改common.py
- 4.2 修改yolo.py
- 4.3 修改train_dual.py
- 五、yaml模型文件
- 5.1 模型改进
- 六、成功运行结果
二、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 修改common.py
此处需要修改的文件是models/common.py
common.py中定义了网络结构的通用模块，我们想要加入新的模块就只需要将模块代码放到这个文件内即可。
HWD的实现过程中使用的pytorch_wavelets包需要自行安装：
pip install pytorch_wavelets
HWD模块添加后如下：
注意❗：在4.2小节中的yolo.py文件中需要声明的模块名称为：HWD。
4.2 修改yolo.py
此处需要修改的文件是models/yolo.py
yolo.py用于函数调用，我们只需要将common.py中定义的新的模块名添加到parse_model函数下即可。
HWD模块添加后如下：
还需在此函数下添加如下代码：
elif m in (HWD,):
args = [ch[f], ch[f]]
4.3 修改train_dual.py
在train_dual.py文件的第314行关闭amp，将其设置为False。
with torch.cuda.amp.autocast(False):
pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
五、yaml模型文件
5.1 模型改进
在代码配置完成后，配置模型的YAML文件。
此处以models/detect/yolov9-c.yaml为例，在同目录下创建一个用于自己数据集训练的模型文件yolov9-c-hwd.yaml。
将yolov9-c.yaml中的内容复制到yolov9-c-hwd.yaml文件下，修改nc数量等于自己数据中目标的数量。
📌 修改方法是将HWD模块替换YOLOv9网络中的ADown模块。HWD受无损信息变换方法的启发，引入Haar小波变换到下采样模块中，旨在尽可能地保留图像信息，使改进后的模型在下采样过程中能够提取更具判别性的特征，从而提高模型性能。
# YOLOv9
# parameters
nc: 1 # number of classes
depth_multiple: 1.0 # model depth multiple
width_multiple: 1.0 # layer channel multiple
#activation: nn.LeakyReLU(0.1)
#activation: nn.ReLU()
# anchors
anchors: 3
# YOLOv9 backbone
backbone:
[
[-1, 1, Silence, []],
# conv down
[-1, 1, Conv, [64, 3, 2]], # 1-P1/2
# conv down
[-1, 1, Conv, [128, 3, 2]], # 2-P2/4
# elan-1 block
[-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 3
# avg-conv down
[-1, 1, HWD, [256]], # 4-P3/8
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 5
# avg-conv down
[-1, 1, HWD, [512]], # 6-P4/16
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 7
# avg-conv down
[-1, 1, HWD, [512]], # 8-P5/32
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 9
]
# YOLOv9 head
head:
[
# elan-spp block
[-1, 1, SPPELAN, [512, 256]], # 10
# up-concat merge
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 7], 1, Concat, [1]], # cat backbone P4
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 13
# up-concat merge
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 5], 1, Concat, [1]], # cat backbone P3
# elan-2 block
[-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small)
# avg-conv-down merge
[-1, 1, ADown, [256]],
[[-1, 13], 1, Concat, [1]], # cat head P4
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium)
# avg-conv-down merge
[-1, 1, ADown, [512]],
[[-1, 10], 1, Concat, [1]], # cat head P5
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 22 (P5/32-large)
# multi-level reversible auxiliary branch
# routing
[5, 1, CBLinear, [[256]]], # 23
[7, 1, CBLinear, [[256, 512]]], # 24
[9, 1, CBLinear, [[256, 512, 512]]], # 25
# conv down
[0, 1, Conv, [64, 3, 2]], # 26-P1/2
# conv down
[-1, 1, Conv, [128, 3, 2]], # 27-P2/4
# elan-1 block
[-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 28
# avg-conv down fuse
[-1, 1, ADown, [256]], # 29-P3/8
[[23, 24, 25, -1], 1, CBFuse, [[0, 0, 0]]], # 30
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 31
# avg-conv down fuse
[-1, 1, ADown, [512]], # 32-P4/16
[[24, 25, -1], 1, CBFuse, [[1, 1]]], # 33
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 34
# avg-conv down fuse
[-1, 1, ADown, [512]], # 35-P5/32
[[25, -1], 1, CBFuse, [[2]]], # 36
# elan-2 block
[-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 37
# detection head
# detect
[[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]], # DualDDetect(A3, A4, A5, P3, P4, P5)
]
六、成功运行结果
分别打印网络模型可以看到HWD模块已经加入到模型中，并可以进行训练了。
yolov9-c-hwd：
from n params module arguments
0 -1 1 0 models.common.Silence []
1 -1 1 1856 models.common.Conv [3, 64, 3, 2]
2 -1 1 73984 models.common.Conv [64, 128, 3, 2]
3 -1 1 212864 models.common.RepNCSPELAN4 [128, 256, 128, 64, 1]
4 -1 1 262656 models.common.HWD [256, 256]
5 -1 1 847616 models.common.RepNCSPELAN4 [256, 512, 256, 128, 1]
6 -1 1 1049600 models.common.HWD [512, 512]
7 -1 1 2857472 models.common.RepNCSPELAN4 [512, 512, 512, 256, 1]
8 -1 1 1049600 models.common.HWD [512, 512]
9 -1 1 2857472 models.common.RepNCSPELAN4 [512, 512, 512, 256, 1]
10 -1 1 656896 models.common.SPPELAN [512, 512, 256]
11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
12 [-1, 7] 1 0 models.common.Concat [1]
13 -1 1 3119616 models.common.RepNCSPELAN4 [1024, 512, 512, 256, 1]
14 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
15 [-1, 5] 1 0 models.common.Concat [1]
16 -1 1 912640 models.common.RepNCSPELAN4 [1024, 256, 256, 128, 1]
17 -1 1 164352 models.common.ADown [256, 256]
18 [-1, 13] 1 0 models.common.Concat [1]
19 -1 1 2988544 models.common.RepNCSPELAN4 [768, 512, 512, 256, 1]
20 -1 1 656384 models.common.ADown [512, 512]
21 [-1, 10] 1 0 models.common.Concat [1]
22 -1 1 3119616 models.common.RepNCSPELAN4 [1024, 512, 512, 256, 1]
23 5 1 131328 models.common.CBLinear [512, [256]]
24 7 1 393984 models.common.CBLinear [512, [256, 512]]
25 9 1 656640 models.common.CBLinear [512, [256, 512, 512]]
26 0 1 1856 models.common.Conv [3, 64, 3, 2]
27 -1 1 73984 models.common.Conv [64, 128, 3, 2]
28 -1 1 212864 models.common.RepNCSPELAN4 [128, 256, 128, 64, 1]
29 -1 1 164352 models.common.ADown [256, 256]
30 [23, 24, 25, -1] 1 0 models.common.CBFuse [[0, 0, 0]]
31 -1 1 847616 models.common.RepNCSPELAN4 [256, 512, 256, 128, 1]
32 -1 1 656384 models.common.ADown [512, 512]
33 [24, 25, -1] 1 0 models.common.CBFuse [[1, 1]]
34 -1 1 2857472 models.common.RepNCSPELAN4 [512, 512, 512, 256, 1]
35 -1 1 656384 models.common.ADown [512, 512]
36 [25, -1] 1 0 models.common.CBFuse [[2]]
37 -1 1 2857472 models.common.RepNCSPELAN4 [512, 512, 512, 256, 1]
38[31, 34, 37, 16, 19, 22] 1 21542822 DualDDetect [1, [512, 512, 512, 256, 512, 512]]
yolov9-c-hwd summary: 601 layers, 51583014 parameters, 49258246 gradients, 239.5 GFLOPs
