YOLOv8 改进 006:SPD-Conv 空间深度转换卷积替换 CNN 传统卷积神经网络

论文题目:《No More Strided Convolutions or Pooling: A New CNN Building Block for Low-Resolution Images and Small Objects》

摘       要:卷积神经网络(CNN)在许多计算机视觉任务,如图像分类和对象检测方面取得了巨大成功。然而,它们在处理更困难的任务时,如低分辨率图像或小物体时,性能会迅速下降。这是由于现有 CNN 常见的设计体系结构中有缺陷,即使用卷积步长和/或池化层,这导致了细粒度信息的丢失和较低效的特征表示的学习。为此,我们提出了一个名为 SPD-Conv 的新的 CNN 构建块来代替每个卷积步长和每个池化层(因此完全消除了它们)。SPD-Conv 由一个空间到深度(SPD)层和一个无卷积步长(Conv)层组成,可以应用于大多数 CNN 体系结构。我们从两个最具代表性的计算机视觉任务,即目标检测和图像分类来解释这个新设计。然后,我们将 SPD-Conv 应用于 YOLOv5 和 ResNet,创建了新的 CNN 架构,并通过经验证明,我们的方法明显优于最先进的深度学习模型,特别是在处理低分辨率图像和小物体等更困难的任务时。

论文地址https://arxiv.org/pdf/2208.03641v1

官方源码GitHub - LabSAINT/SPD-Conv

1. SPD-Conv 简介

卷积神经网络(CNN)在处理低分辨率图像和小物体时性能低下,这一问题根源于使用步长卷积和池化层导致的细粒度信息丢失。

SPD-Conv(空间到深度卷积)的基本原理是用于改进传统卷积神经网络(CNN)中对低分辨率图像和小物体处理的性能。它主要通过以下几个关键步骤实现:

(1)替换步长卷积和池化层:SPD-Conv 被设计用来替代传统 CNN 架构中的步长卷积层和池化层。步长卷积和池化层在处理低分辨率图像或小物体时会导致细粒度信息的丢失。

(2)空间到深度(SPD)层:一个转换层,将输入图像的空间维度转换为深度维度,从而在不丢失信息的情况下增加特征图的深度。之所以使用 SPD 层,是因为在处理低分辨率图像和小物体时,需要保留尽可能多的空间信息。SPD 层通过将空间维度的信息转换为深度维度,避免了传统步长卷积和池化操作中的信息丢失。

(3)非步长卷积层:在 SPD 转换后应用的卷积层,不使用步长,以保留细粒度信息。在 SPD 层之后应用非步长(即步长为1)卷积层,是因为非步长卷积能够在不减少特征图尺寸的情况下进行特征提取,进一步保持了图像的细粒度信息,这对于提高低分辨率图像和小物体的识别性能至关重要。

假设我们有一个低分辨率的图像,其中包含几个小的物体,我们需要对这些物体进行识别和分类。

在传统的 CNN 架构中,如果我们直接应用步长卷积和池化层,随着网络层次的加深,图像的空间分辨率会逐渐减少,导致小物体的细节信息丢失,从而使得网络难以准确识别这些小物体。

使用 SPD-Conv 代替步长卷积和池化层

(1)空间到深度(SPD)层的应用:

  • 初始特征图尺寸为(32x32),包含小物体的细节信息。
  • 应用 SPD 层后,假设将空间分辨率降低为(16x16)),同时将这部分减少的空间信息转移到通道维度上,从而通道数增加但没有信息丢失。

(2)非步长卷积层的应用:

  • 经过SPD层处理后,特征图的尺寸变为(16 x16),但通道数增加,假设从原来的 64 通道增加到 256 通道。
  • 在这个增加通道数的特征图上应用非步长卷积层,这个卷积层不减少特征图的空间尺寸,而是通过学习这些增加的通道中的信息来提取重要特征。
  • 这样,即使是在较小的空间分辨率上,模型也能有效捕捉到小物体的细节信息。

这种结合使用 SPD 层和非步长卷积层的方法,使得 CNN 能够更好地处理小物体和低分辨率图像中的挑战,提高了模型在这些复杂场景下的性能和鲁棒性。

以下是我对这张图的理解

(1)图(a)显示了标准的特征图,具有通道数 C1、宽度和高度;

(2)图(b)展示了空间到深度操作,其中空间信息被重新排列到深度通道;

(3)图(c)显示了结果特征图的深度增加;

(4)图(d)表示在 SPD 操作之后应用的非步长卷积层;

(5)图(e)展示了经过步长为 1 的卷积后的输出特征图,该特征图保持了空间分辨率但改变了深度维度。

2. 项目环境

  • 解释器:3.9.19
  • 框架:Pytorch 2.0.0 + CUDA 11.8
  • 系统:Win10 / Ubuntu 20.04

3. 核心代码

以下提供两份代码供大家学习,添加教程使用代码 1。

代码 1:

import torch
import torch.nn as nn

__all__ = ['SPDConv']

class SPDConv(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)

代码 2:

import torch
import torch.nn as nn
 
__all__ = ['SPDConv']
 
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 SPDConv(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__()
        c1 = c1 * 4
        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):
        x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
        """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."""
        x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
        return self.act(self.conv(x))

4. 添加方法

 第 1 步 :在 ultralytics/nn/add_modules/ 目录下新建 Python 源文件 SPDConv.py,将以上 SPD-Conv 空间深度转换卷积的核心代码(代码 1)复制粘贴至 SPDConv.py 文件中。

 第 2 步 :定位到 ultralytics/nn/add_modules/ 目录下的 __init__.py 文件,加入 SPDConv。

from .SPDConv import *

 第 3 步 :定位到 ultralytics/nn/ 目录下的 tasks.py 文件,找到 parse_model 函数添加以下代码。

# -------------SPD-Conv------------
elif m is SPDConv:
    c2 = 4 * ch[f]
# ---------------END---------------

添加完成之后,需导入 SPDConv 模块。

from ultralytics.nn.add_modules import *

  第 4 步 :在 ultralytics\cfg\models\add\ 目录下新建 YAML 文件 yolov8-SPDConv.yaml,复制 yolov8.yaml 中的代码粘贴至此处,使用 SPD-Conv 空间深度转换卷积替换 CNN 传统卷积神经网络的下采样。

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect

# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
  s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
  m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
  l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
  x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs


# YOLOv8.0n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Focus, [64, 3]]       # 0-P1/2

  - [-1, 1, Conv, [128, 3, 1]]    # 1-P2/4
  - [-1, 1, SPDConv, [1]]         # 2-P2/4
  - [-1, 3, C2f, [128, True]]     # 3

  - [-1, 1, Conv, [256, 3, 1]]    # 4-P3/8
  - [-1, 1, SPDConv, [1]]         # 5
  - [-1, 6, C2f, [256, True]]     # 6

  - [-1, 1, Conv, [512, 3, 1]]    # 7-P4/16
  - [-1, 1, SPDConv, [1]]         # 8
  - [-1, 6, C2f, [512, True]]     # 9

  - [-1, 1, Conv, [1024, 3, 1]]   # 10-P5/32
  - [-1, 1, SPDConv, [1]]         # 11
  - [-1, 3, C2f, [1024, True]]    # 12
  - [-1, 1, SPPF, [1024, 5]]      # 13

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 8], 1, Concat, [1]] # cat backbone P4
  - [-1, 3, C2f, [512]] # 16

  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 5], 1, Concat, [1]] # cat backbone P3
  - [-1, 3, C2f, [256]] # 19 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 1]] # 20
  - [-1, 1, SPDConv, [1]] # 21
  - [[-1, 16], 1, Concat, [1]] # cat head P4
  - [-1, 3, C2f, [512]] # 23 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 1]] # 24
  - [-1, 1, SPDConv, [1]] # 25
  - [[-1, 13], 1, Concat, [1]] # cat head P5
  - [-1, 3, C2f, [1024]] # 27 (P5/32-large)

  - [[19, 23, 27], 1, Detect, [nc]] # Detect(P3, P4, P5)

5. 训练代码

在 YOLOv8.2 项目的主目录下新建 Python 源文件 train.py,将以下代码复制粘贴至该文件中。

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':

    model = YOLO(r'D:\Lab\YOLOv8.2\ultralytics\cfg\add\yolov8s-SPDConv.yaml')

    # model.load('yolov8n.pt') # 是否加载预训练权重,科研不建加载否则很难提升精度

    model.train(
        data=r'The YAML file address of your own dataset.',
        cache=False,
        imgsz=640,
        epochs=300,
        single_cls=False,  # 是否是单类别检测
        batch=2,
        close_mosaic=0,
        workers=0,
        device='0',
        optimizer='SGD',   # using SGD
        # resume='runs/train/exp/weights/last.pt',   # 如过想续训就设置 last.pt 的地址
        amp=False,                                   # 如果出现训练损失为 Nan 可以关闭 amp
        project='runs/train',
        name='exp',
    )

6. 训练过程

右键直接运行 train.py 就可以开始训练模型,代码首先会输出模型的基本信息(模型有几个卷积层、池化层、全连接层构成)和运行的记录,如下图所示。

之后将持续发布 “小目标检测头” 和 “SPD-Conv 空间深度转换卷积” 结合改进机制!

 欢迎大家订阅我的专栏一起学习 YOLO!(o^^o)

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值