主干网络篇 | YOLOv8更换主干网络之ShuffleNetV2

前言：Hello大家好，我是小哥谈。ShuffleNetV2是一种轻量级的神经网络架构，用于图像分类和目标检测任务。它是ShuffleNet的改进版本，旨在提高模型的性能和效率。ShuffleNetV2相比于之前的版本，在保持模型轻量化的同时，提高了模型的准确性和性能。它在计算资源有限的设备上具有较好的应用潜力！~🌈

🚀1. 基础概念

🚀2.网络结构

🚀3.添加步骤

🚀4.改进方法

🍀🍀步骤1：block.py文件修改

🍀🍀步骤2：__init__.py文件修改

🍀🍀步骤3：tasks.py文件修改

🍀🍀步骤4：创建自定义yaml文件

🍀🍀步骤5：新建train.py文件

🍀🍀步骤6：模型训练测试

🚀1. 基础概念

ShuffleNetV2是一种轻量级的神经网络架构，用于图像分类和目标检测任务。它是ShuffleNet的改进版本，旨在提高模型的性能和效率。

ShuffleNetV2的主要特点包括：

分组卷积：通过将输入通道分成多个组，并在组内进行卷积操作，减少了计算量和参数数量。
逐点卷积：使用1x1的卷积核进行逐点卷积，用于调整通道数和特征图的维度。
通道重排：通过将输入特征图按通道进行重排，实现信息的混洗和交互，增强了特征的表达能力。
瓶颈结构：采用瓶颈结构，即先降维再升维，减少了计算量和参数数量。
网络设计：ShuffleNet V2通过堆叠多个ShuffleNet单元来构建整个网络，可以根据任务的需求进行不同层数和宽度的配置。

ShuffleNetV2相比于之前的版本，在保持模型轻量化的同时，提高了模型的准确性和性能。它在计算资源有限的设备上具有较好的应用潜力。

shuffleNetV2这篇论文比较硬核，提出了不少新的思想，推荐大家可以看看论文原文。主要思想包括：

模型的计算复杂度不能只看FLOPs,还需要参考一些其他的指标
作者提出了4条如何设计高效网络的准则
基于该准则提出了新的block设置

FLOPS网上有两种：FLOPS和 FLOPs

FLOPS：全大写，指每秒浮点运算次数，可以理解为计算的速度，是衡量硬件性能的一个指标 (硬件)
FLOPs：s小写，指浮点运算数，理解为计算量，可以用来衡量算法/模型的复杂度，（模型）在论文中常用GFLOPs(1 GFLOPs = 10^9FLOPs)

ShuffleNetV2网络结构：

原理图：

其中，a、b为ShuffleNetV1原理图，c、d为ShuffleNetV2原理图。

论文题目：《ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design》

论文地址： https://arxiv.org/pdf/1807.11164.pdf

代码实现： GitHub - megvii-model/ShuffleNet-Series

🚀2.网络结构

本文的改进是基于YOLOv8，关于其网络结构具体如下图所示：

YOLOv8官方仓库地址：

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite

针对本文的改进，作者将所使用的含有预训练权重文件的YOLOv8完整源码进行了上传，大家可在我的“资源”中自行下载。

🚀3.添加步骤

针对本文的改进，具体步骤如下所示：👇

步骤1：block.py文件修改

步骤2：__init__.py文件修改

步骤3：tasks.py文件修改

步骤4：创建自定义yaml文件

步骤5：新建train.py文件

步骤6：模型训练测试

🚀4.改进方法

🍀🍀步骤1：block.py文件修改

在源码中找到block.py文件，具体位置是ultralytics/nn/modules/block.py，然后将ShuffleNetV2模块代码添加到block.py文件末尾位置。

ShuffleNetV2模块代码：

# ShuffleNetv2核心代码
# By CSDN 小哥谈
import torch
import torch.nn as nn

def channel_shuffle(x, groups):
    batchsize, num_channels, height, width = x.data.size()
    channels_per_group = num_channels // groups
    x = x.view(batchsize, groups, channels_per_group, height, width)
    x = torch.transpose(x, 1, 2).contiguous()
    x = x.view(batchsize, -1, height, width)
    return x


class CBRM(nn.Module):  # Conv BN ReLU Maxpool2d
    def __init__(self, c1, c2):
        super(CBRM, 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, ch_in, ch_out, stride):
        super(Shuffle_Block, self).__init__()

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

        branch_features = ch_out // 2
        assert (self.stride != 1) or (ch_in == branch_features << 1)

        if self.stride > 1:
            self.branch1 = nn.Sequential(
                self.depthwise_conv(ch_in, ch_in, kernel_size=3, stride=self.stride, padding=1),
                nn.BatchNorm2d(ch_in),

                nn.Conv2d(ch_in, 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(ch_in 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)
            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

再然后，在block.py文件最上方下图所示位置加入CBRM、Shuffle_Block。

🍀🍀步骤2：init.py文件修改

在源码中找到__init__.py文件，具体位置是ultralytics/nn/modules/__init__.py。

修改1：加入CBRM、Shuffle_Block，具体如下图所示：

修改2：加入CBRM、Shuffle_Block，具体如下图所示：

🍀🍀步骤3：tasks.py文件修改

在源码中找到tasks.py文件，具体位置是ultralytics/nn/tasks.py。

修改1：在下图所示位置导入类名CBRM、Shuffle_Block。

修改2：找到parse_model函数（736行左右），在下图中所示位置添加如下代码。

 # -------ShuffleNetv2------------
        elif m in [CBRM, Shuffle_Block]:
            c1, c2 = ch[f], args[0]
            if c2 != nc:
                c2 = make_divisible(min(c2, max_channels) * width, 8)
            args = [c1, c2, *args[1:]]
        # --------------------------------

具体添加位置如下图所示：

🍀🍀步骤4：创建自定义yaml文件

在源码ultralytics/cfg/models/v8目录下创建yaml文件，并命名为：yolov8_ShuffleNetV2.yaml。具体如下图所示：

yolov8_ShuffleNetV2.yaml文件完整代码如下所示：

# 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, CBRM, [ 32 ] ] # 0-P2/4
  - [ -1, 1, Shuffle_Block, [ 128, 2 ] ]  # 1-P3/8
  - [ -1, 3, Shuffle_Block, [ 128, 1 ] ]  # 2
  - [ -1, 1, Shuffle_Block, [ 256, 2 ] ]  # 3-P4/16
  - [ -1, 7, Shuffle_Block, [ 256, 1 ] ]  # 4
  - [ -1, 1, Shuffle_Block, [ 512, 2 ] ]  # 5-P5/32
  - [ -1, 3, Shuffle_Block, [ 512, 1 ] ]  # 6


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

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

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 9], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C2f, [512]]  # 15 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 6], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C2f, [1024]]  # 18 (P5/32-large)

  - [[12, 15, 18], 1, Detect, [nc]]  # Detect(P3, P4, P5)

🍀🍀步骤5：新建train.py文件

在源码根目录下新建train.py文件，文件完整代码如下所示：

from ultralytics import YOLO

# Load a model
model = YOLO(r'C:\Users\Lenovo\PycharmProjects\ultralytics-main\ultralytics\cfg\models\v8\yolov8_ShuffleNetV2.yaml')  # build a new model from YAML
model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
model = YOLO(r'C:\Users\Lenovo\PycharmProjects\ultralytics-main\ultralytics\cfg\models\v8\yolov8_ShuffleNetV2.yaml').load('yolov8n.pt')  # build from YAML and transfer weights

# Train the model
model.train(data=r'C:\Users\Lenovo\PycharmProjects\ultralytics-main\ultralytics\cfg\datasets\helmet.yaml', epochs=100, imgsz=640)