SPD模块
理论知识
问题背景
在图像分辨率低或物体很小的复杂任务中,传统CNN架构的性能会急速下降。这是因为传统CNN架构中常见的跨行卷积和池化层会导致细粒度信息的丢失和学习不太有效的特征表示。在早期CNN架构中,这种设计没有表现出明显的缺点。因为,大多数被研究的场景都是“友好的”,其中图像具有良好的分辨率,对象大小适中;因此,有大量的冗余像素信息,跨卷积和池化可以方便地跳过,模型仍然可以很好地学习特征。然而,在图像模糊或物体很小的复杂任务中,对冗余信息的大量假设不再成立,当前的设计开始遭受细粒度信息丢失和特征学习不良的影响。
SPD-Conv结构
由结构图可以看出SPD-Conv的结构十分简单,就是一个空间到深度(SPD)层,后面跟着一个非跨行卷积层(其实就是stride为1且kernal_size=1的卷积)。SPD层对特征映射X进行下采样,但保留通道维度中的所有信息,因此没有信息丢失。受到图像变换技术的启发,该技术在将原始图像馈送到神经网络之前重新缩放原始图像,但我们基本上将其推广到内部和整个网络中的降采样特征映射。此外,在每个SPD之后添加了一个非跨行卷积操作,以使用添加的卷积层中的可学习参数减少(增加)通道数。
在YOLO中的应用
在YOLO中的结构图
论文中在YOLOv5中实践了,其结构图如下图所示:
代码
在common.py
中添加如下代码:
class space_to_depth(nn.Module):
# Changing the dimension of the Tensor
def __init__(self, dimension=1):
super().__init__()
self.d = dimension
def forward(self, x):
return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
这段代码定义了一个名为space_to_depth
的自定义PyTorch
模块(nn.Module
的子类),用于改变张量的维度。
def __init__(self, dimension=1):
是构造函数,它接受一个可选的参数dimension
,用于指定要改变的张量维度,默认为1。构造函数通过调用super().__init__()
来初始化nn.Module
的基类。def forward(self, x):
是模块的前向传播函数,它接受一个张量x作为输入,并返回该模块的输出。x[..., ::2, ::2]
表示选取张量x中特定索引位置的元素,…表示在所有非指定维度上进行选取,::2
表示以步长为2进行选取。这里是选取张量x中索引位置为偶数的行和列。x[..., 1::2, ::2]
是选取索引位置为奇数的行和偶数的列。x[..., ::2, 1::2]
是选取索引位置为偶数的行和奇数的列。x[..., 1::2, 1::2]
是选取索引位置为奇数的行和列。
torch.cat()
函数用于将选取的四个部分按照维度1进行连接,得到最终的输出张量。
通过这样的方式,forward
函数将输入张量x重新排列,并返回一个新的张量作为模块的输出。这个操作通常用于空间到深度的转换,其中输入张量被重组为更深的张量,维度之间的空间信息被整合在一起。
在yolo.py
中添加
elif m is Contract:
c2 = ch[f] * args[0] ** 2
elif m is Expand:
c2 = ch[f] // args[0] ** 2
elif m is space_to_depth:
c2 = 4*ch[f]
else:
c2 = ch[f]
在yaml
文件中添加:
# Parameters
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 1]], # 1
[-1,1,space_to_depth,[1]], # 2 -P2/4
[-1, 3, C3, [128]], # 3
[-1, 1, Conv, [256, 3, 1]], # 4
[-1,1,space_to_depth,[1]], # 5 -P3/8
[-1, 6, C3, [256]], # 6
[-1, 1, Conv, [512, 3, 1]], # 7-P4/16
[-1,1,space_to_depth,[1]], # 8 -P4/16
[-1, 9, C3, [512]], # 9
[-1, 1, Conv, [1024, 3, 1]], # 10-P5/32
[-1,1,space_to_depth,[1]], # 11 -P5/32
[-1, 3, C3, [1024]], # 12
[-1, 1, SPPF, [1024, 5]], # 13
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]], # 14
[-1, 1, nn.Upsample, [None, 2, 'nearest']], # 15
[[-1, 9], 1, Concat, [1]], # 16 cat backbone P4
[-1, 3, C3, [512, False]], # 17
[-1, 1, Conv, [256, 1, 1]], # 18
[-1, 1, nn.Upsample, [None, 2, 'nearest']], # 19
[[-1, 6], 1, Concat, [1]], # 20 cat backbone P3
[-1, 3, C3, [256, False]], # 21 (P3/8-small)
[-1, 1, Conv, [256, 3, 1]], # 22
[-1,1,space_to_depth,[1]], # 23 -P2/4
[[-1, 18], 1, Concat, [1]], # 24 cat head P4
[-1, 3, C3, [512, False]], # 25 (P4/16-medium)
[-1, 1, Conv, [512, 3, 1]], # 26
[-1,1,space_to_depth,[1]], # 27 -P2/4
[[-1, 14], 1, Concat, [1]], # 28 cat head P5
[-1, 3, C3, [1024, False]], # 29 (P5/32-large)
[[21, 25, 29], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
实验
VisDrone数据集
VisDrone数据集是由天津大学开发的一个无人机视角目标检测数据集,其中包括多个类别,不同尺寸大小的目标,既有大中尺寸的目标,如:车辆等;也有微小目标,如:行人。
VisDrone数据集地址
实验结果
以YOLOv5s作为基准进行对比,VisDrone数据集在YOLOv5s上的mAP为32.9%。通过修改主干的卷积为SPD,精度有较高的提升,mAP达到34.3%。证明了该模块在小目标检测上的有效性。
相关模型以及训练结果地址:百度云:c60y