YOLOv9改进策略【模型轻量化】| ShufflenetV2，通过通道划分构建高效网络

一、本文介绍

本文记录的是基于ShufflenetV2的YOLOv9目标检测轻量化改进方法研究。FLOPs是评价模型复杂独的重要指标，但其无法考虑到模型的内存访问成本和并行度，因此本文在YOLOv9的基础上引入ShufflenetV2，使其在在保持准确性的同时提高模型的运行效率。

模型	参数量	计算量	推理速度(bs=32)
YOLOv9-c	50.69M	236.6GFLOPs	32.1ms
Improved	42.88M	194.5GFLOPs	23.2ms

文章目录

一、本文介绍
二、ShuffleNet V2设计原理
三、ShuffleNet V2基础模块的实现代码
四、添加步骤
- 4.1 修改common.py
- 4.2 修改yolo.py
五、yaml模型文件
- 5.1 模型改进⭐
六、成功运行结果

二、ShuffleNet V2设计原理

ShuffleNet V2是一种高效的卷积神经网络架构，其模型结构及优势如下：

模型结构：
- 回顾ShuffleNet v1：ShuffleNet是一种广泛应用于低端设备的先进网络架构，为增加在给定计算预算下的特征通道数量，采用了点组卷积和瓶颈结构，但这增加了内存访问成本（MAC），且过多的组卷积和元素级“Add”操作也存在问题。
- 引入Channel Split和ShuffleNet V2：为解决上述问题，引入了名为Channel Split的简单操作。在每个单元开始时，将 $c$ 个特征通道的输入分为两个分支，分别具有 $c - c^{'}$ 和 $c^{'}$ 个通道。一个分支保持不变，另一个分支由三个具有相同输入和输出通道的卷积组成，以满足G1（平衡卷积，即相等的通道宽度可最小化MAC）。两个 $\times 1$ 卷积不再是组式的，这部分是为了遵循G2（避免过多的组卷积增加MAC），部分是因为拆分操作已经产生了两个组。卷积后，两个分支连接，通道数量保持不变，并使用与ShuffleNet v1相同的“通道洗牌”操作来实现信息通信。对于空间下采样，单元进行了略微修改，删除了通道拆分操作，使输出通道数量加倍。
- 整体网络结构：通过反复堆叠构建块来构建整个网络，设置 $c^{'} = c /2$ ，整体网络结构与ShuffleNet v1相似，并在全局平均池化之前添加了一个额外的 $\times 1$ 卷积层来混合特征。
优势：
- 高效且准确：遵循了高效网络设计的所有准则，每个构建块的高效率使其能够使用更多的特征通道和更大的网络容量，并且在每个块中，一半的特征通道直接通过块并加入下一个块，实现了一种特征重用模式，类似于DenseNet，但更高效。
- 速度优势明显：在与其他网络架构的比较中，ShuffleNet v2在速度方面表现出色，特别是在GPU上明显快于其他网络（如MobileNet v2、ShuffleNet v1和Xception）。在ARM上，ShuffleNet v1、Xception和ShuffleNet v2的速度相当，但MobileNet v2较慢，这是因为MobileNet v2的MAC较高。
- 兼容性好：可以与其他技术（如Squeeze - and - excitation模块）结合进一步提高性能。

论文：https://arxiv.org/pdf/1807.11164.pdf
源码：https://gitcode.com/gh_mirrors/sh/ShuffleNet-Series/blob/master/ShuffleNetV2/blocks.py?utm_source=csdn_github_accelerator&isLogin=1

三、ShuffleNet V2基础模块的实现代码

ShuffleNet V2基础模块的实现代码如下：

def channel_shuffle(x, groups):
    batchsize, num_channels, height, width = x.data.size()
    channels_per_group = num_channels // groups

    # reshape
    x = x.view(batchsize, groups,
               channels_per_group, height, width)

    x = torch.transpose(x, 1, 2).contiguous()

    # flatten
    x = x.view(batchsize, -1, height, width)

    return x


class conv_bn_relu_maxpool(nn.Module):
    def __init__(self, c1, c2):  # ch_in, ch_out
        super(conv_bn_relu_maxpool, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(c1, c2, kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(c2),
            nn.ReLU(inplace=True),
        )
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

    def forward(self, x):
        return self.maxpool(self.conv(x))


class Shuffle_Block(nn.Module):
    def __init__(self, inp, oup, stride):
        super(Shuffle_Block, self).__init__()

        if not (1 <= stride <= 3):
            raise ValueError('illegal stride value')
        self.stride = stride

        branch_features = oup // 2
        assert (self.stride != 1) or (inp == branch_features << 1)

        if self.stride > 1:
            self.branch1 = nn.Sequential(
                self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1),
                nn.BatchNorm2d(inp),
                nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(branch_features),
                nn.ReLU(inplace=True),
            )

        self.branch2 = nn.Sequential(
            nn.Conv2d(inp if (self.stride > 1) else branch_features,
                      branch_features, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(branch_features),
            nn.ReLU(inplace=True),
            self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
            nn.BatchNorm2d(branch_features),
            nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(branch_features),
            nn.ReLU(inplace=True),
        )

    @staticmethod
    def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
        return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)

    def forward(self, x):
        if self.stride == 1:
            x1, x2 = x.chunk(2, dim=1)  # 按照维度1进行split
            out = torch.cat((x1, self.branch2(x2)), dim=1)
        else:
            out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)

        out = channel_shuffle(out, 2)

        return out

四、添加步骤

4.1 修改common.py

此处需要修改的文件是models/common.py

common.py中定义了网络结构的通用模块，我们想要加入新的模块就只需要将模块代码放到这个文件内即可。

此时需要将上方实现的代码添加到common.py中。

在这里插入图片描述

注意❗：在4.2小节中的yolo.py文件中需要声明的模块名称为：conv_bn_relu_maxpool和Shuffle_Block。

4.2 修改yolo.py

此处需要修改的文件是models/yolo.py

yolo.py用于函数调用，我们只需要将common.py中定义的新的模块名添加到parse_model函数下即可。

conv_bn_relu_maxpool模块以及Shuffle_Block模块添加后如下：

在这里插入图片描述

五、yaml模型文件

5.1 模型改进⭐

在代码配置完成后，配置模型的YAML文件。

此处以models/detect/yolov9-c.yaml为例，在同目录下创建一个用于自己数据集训练的模型文件yolov9-c-shufflenetv2.yaml。

将yolov9-c.yaml中的内容复制到yolov9-c-shufflenetv2.yaml文件下，修改nc数量等于自己数据中目标的数量。

📌 模型的修改方法是将YOLOv9的骨干网络替换成Shufflenet V2，ShuffleNet V2 在设计上注重减少内存访问成本并提高并行度，这有助于在保持准确性的同时提高模型的运行效率。相比YOLOv9原骨干网络，ShuffleNet V2 具有更低的计算复杂度，能够在相同或更少的计算资源下完成推理，对于实时性要求较高的任务具有重要意义。

结构如下：

# 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_bn_relu_maxpool, [64, 3, 2]],  # 1-P1/2

   # conv down
   [-1, 1, Shuffle_Block, [ 128, 2 ]],  # 2-P2/4
   [-1, 3, Shuffle_Block, [ 128, 1 ]],  # 3

   [-1, 1, Shuffle_Block, [ 256, 2 ]],  # 4-P4/16 
   [-1, 7, Shuffle_Block, [ 256, 1 ]],  # 5

   [-1, 1, Shuffle_Block, [ 512, 2 ]],  # 6-P4/16
   [-1, 3, Shuffle_Block, [ 512, 1 ]],  # 7
   
  ]

# YOLOv9 head
head:
  [
   # elan-spp block
   [-1, 1, SPPELAN, [512, 256]],  # 10

   # up-concat merge
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 5], 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, 3], 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, 11], 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, 8], 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
   [3, 1, CBLinear, [[256]]], # 23
   [5, 1, CBLinear, [[256, 512]]], # 24
   [7, 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
   [[21, 22, 23, -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
   [[22, 23, -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
   [[23, -1], 1, CBFuse, [[2]]], # 36

   # elan-2 block
   [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 37
   
   
   
   # detection head

   # detect
   [[29, 32, 35, 14, 17, 20], 1, DualDDetect, [nc]],  # DualDDetect(A3, A4, A5, P3, P4, P5)
  ]

六、成功运行结果

分别打印网络模型可以看到Shuffle_Block已经加入到模型中，并可以进行训练了。

yolov9-c-shufflenetv2：

                 from  n    params  module                                  arguments                     
  0                -1  1         0  models.common.Silence                   []                            
  1                -1  1      1856  models.common.conv_bn_relu_maxpool      [3, 64]                       
  2                -1  1     14080  models.common.Shuffle_Block             [64, 128, 2]                  
  3                -1  3     27456  models.common.Shuffle_Block             [128, 128, 1]                 
  4                -1  1     52736  models.common.Shuffle_Block             [128, 256, 2]                 
  5                -1  7    242816  models.common.Shuffle_Block             [256, 256, 1]                 
  6                -1  1    203776  models.common.Shuffle_Block             [256, 512, 2]                 
  7                -1  3    404736  models.common.Shuffle_Block             [512, 512, 1]                 
  8                -1  1    656896  models.common.SPPELAN                   [512, 512, 256]               
  9                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 10           [-1, 5]  1         0  models.common.Concat                    [1]                           
 11                -1  1   2988544  models.common.RepNCSPELAN4              [768, 512, 512, 256, 1]       
 12                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 13           [-1, 3]  1         0  models.common.Concat                    [1]                           
 14                -1  1    814336  models.common.RepNCSPELAN4              [640, 256, 256, 128, 1]       
 15                -1  1    164352  models.common.ADown                     [256, 256]                    
 16          [-1, 11]  1         0  models.common.Concat                    [1]                           
 17                -1  1   2988544  models.common.RepNCSPELAN4              [768, 512, 512, 256, 1]       
 18                -1  1    656384  models.common.ADown                     [512, 512]                    
 19           [-1, 8]  1         0  models.common.Concat                    [1]                           
 20                -1  1   3119616  models.common.RepNCSPELAN4              [1024, 512, 512, 256, 1]      
 21                 3  1     33024  models.common.CBLinear                  [128, [256]]                  
 22                 5  1    197376  models.common.CBLinear                  [256, [256, 512]]             
 23                 7  1    656640  models.common.CBLinear                  [512, [256, 512, 512]]        
 24                 0  1      1856  models.common.Conv                      [3, 64, 3, 2]                 
 25                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               
 26                -1  1    212864  models.common.RepNCSPELAN4              [128, 256, 128, 64, 1]        
 27                -1  1    164352  models.common.ADown                     [256, 256]                    
 28  [21, 22, 23, -1]  1         0  models.common.CBFuse                    [[0, 0, 0]]                   
 29                -1  1    847616  models.common.RepNCSPELAN4              [256, 512, 256, 128, 1]       
 30                -1  1    656384  models.common.ADown                     [512, 512]                    
 31      [22, 23, -1]  1         0  models.common.CBFuse                    [[1, 1]]                      
 32                -1  1   2857472  models.common.RepNCSPELAN4              [512, 512, 512, 256, 1]       
 33                -1  1    656384  models.common.ADown                     [512, 512]                    
 34          [23, -1]  1         0  models.common.CBFuse                    [[2]]                         
 35                -1  1   2857472  models.common.RepNCSPELAN4              [512, 512, 512, 256, 1]       
 36[29, 32, 35, 14, 17, 20]  1  21542822  models.yolo.DualDDetect                 [1, [512, 512, 512, 256, 512, 512]]
yolov9-c-shufflenetv2 summary: 870 layers, 43094374 parameters, 43094342 gradients, 195.9 GFLOPs