目标检测 YOLOv5 自定义网络结构（YOLOv5-ShuffleNetV2）

flyfish

版本：YOLOv5:v5

已借用的特定自定义网络结构包括：

1. YOLOv5-MobileNetV3

MobileNetV3 Large 
MobileNetV3 Small

1. YOLOv5-ShuffleNetV2

ShuffleNetV2
ShuffleNetV2-Focus
ShuffleNetV2-stem（Pelee的stem模块）

源代码下载地址
训练期间可以使用命令

python train.py  --batch 64 --epochs 300 --data data/coco128.yaml --cfg  models/yolov5-mobilenetv3small.yaml 

yolov5-mobilenetv3small.yaml 可以更换为如下配置

├── yolov5-mobilenetv3large.yaml
├── yolov5-mobilenetv3small.yaml
├── yolov5-shufflenetv2-focus.yaml
├── yolov5-shufflenetv2-stem.yaml
├── yolov5-shufflenetv2.yaml

一种方法是将原始结构组织成
backbones
necks
heads
detectors
losses
utils
例如backbones里面有shufflenetv2，mobilenetv3.py等

其他方式
采用YOLOv5的配置方式
添加模块以根据配置文件生成网络结构。

这里是第二种方式

本文以将ShuffleNetV2的InvertedResidual模块加入到YOLOv5的Backbone中为例说明如何自定义网络结构，加入自定义模块

第一步：添加模块代码

在YOLOv5的源码路径models/common.py里，实现ShuffleNetV2的InvertedResidual
不需要自己手敲代码，直接使用PyTorch官方实现的ShuffleNetV2代码
vision/torchvision/models/shufflenetv2.py

在common.py的顶部加入导入

from torch import Tensor
from typing import Callable, Any, List

将InvertedResidual类和InvertedResidual类需要的channel_shuffle函数都加入到common.py的底部

def channel_shuffle(x: Tensor, groups: int) -> Tensor:
    batchsize, num_channels, height, width = x.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 InvertedResidual(nn.Module):
    def __init__(
        self,
        inp: int,
        oup: int,
        stride: int
    ) -> None:
        super(InvertedResidual, 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),
            )
        else:
            self.branch1 = nn.Sequential()

        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: int,
        o: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        bias: bool = False
    ) -> nn.Conv2d:
        return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)

    def forward(self, x: Tensor) -> Tensor:
        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

第二步 更改解析模块，告诉YOLOv5，我们加入了InvertedResidual模块

目录在models/yolo.py的parse_model函数

if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
            C3, C3TR,InvertedResidual]:

第三步配置

在目录models下新建yolov5-shufflenetv2-focus.yaml文件，配置如下

# parameters
nc: 80  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 0.5  # layer channel multiple

# anchors
anchors:
  - [4,5,  8,10,  13,16]  # P3/8
  - [23,29,  43,55,  73,105]  # P4/16
  - [146,217,  231,300,  335,433]  # P5/32

# Custom backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],    # 0-P2/4
   [-1, 1, InvertedResidual, [128, 2]], # 1-P3/8
   [-1, 3, InvertedResidual, [128, 1]], # 2
   [-1, 1, InvertedResidual, [256, 2]], # 3-P4/16
   [-1, 7, InvertedResidual, [256, 1]], # 4
   [-1, 1, InvertedResidual, [512, 2]], # 5-P5/32
   [-1, 3, InvertedResidual, [512, 1]], # 6
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [128, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P4
   [-1, 1, C3, [128, False]],  # 10

   [-1, 1, Conv, [128, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 2], 1, Concat, [1]],  # cat backbone P3
   [-1, 1, C3, [128, False]],  # 14 (P3/8-small)

   [-1, 1, Conv, [128, 3, 2]],
   [[-1, 11], 1, Concat, [1]],  # cat head P4
   [-1, 1, C3, [128, False]],  # 17 (P4/16-medium)

   [-1, 1, Conv, [128, 3, 2]],
   [[-1, 7], 1, Concat, [1]],  # cat head P5
   [-1, 1, C3, [128, False]],  # 20 (P5/32-large)

   [[14, 17, 20], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

配置参数说明

编写配置文件时，不需要定义输入通道，只需要定义输出通道等参数即可。
输入数据最初是3通道，定义输出通道，该通道也是其他层的输入，
例如

[[-1, 1, conv_bn_hswish,               [16, 2]],                            
 [-1, 1, MobileNetV3_InvertedResidual, [16,  16, 3, 2, 1, 0]],  
 [-1, 1, MobileNetV3_InvertedResidual, [24,  72, 3, 2, 0, 0]], 

代码会自动展开

models.common.conv_bn_hswish                [3, 16, 2]                    
models.common.MobileNetV3_InvertedResidual  [16, 16, 16, 3, 2, 1, 0] 
models.common.MobileNetV3_InvertedResidual  [24, 24, 88, 3, 1, 0, 0] 

conv_bn_hswish的参数
self, c1, c2, stride
含义如下
输入通道是c1=3
输出通道是c2=16
stride=2

MobileNetV3_InvertedResidual的参数

self, inp, oup, hidden_dim, kernel_size, stride, use_se, use_hs

输入通道是inp=16,
输出通道是oup=16,
hidden_dim=16,
kernel_size=3,
stride=2,
use_se=1,
use_hs=0
use_se表示是否使用SELayer
use_hs表示使用h_swish还是ReLU

看过这篇文章的人还看过
目标检测 YOLOv5 – 卷积层和BN层的融合
目标检测 YOLOv5 – Sample Assignment
目标检测 YOLOv5 – 数据增强
目标检测 YOLOv5 – 学习率
目标检测 YOLOv5 – 多机多卡训练
目标检测 YOLOv5 – 浮点取模
目标检测 YOLOv5 – 在多类别中应用NMS（非极大值抑制）
目标检测 YOLOv5 – loss for objectness and classification
目标检测 YOLOv5 – loss for bounding box regression
目标检测 YOLOv5 – 指标计算
目标检测 YOLOv5 – anchor设置
目标检测 YOLOv5 – SPP模块
目标检测 YOLOv5 – 边框预测（bounding box prediction）
目标检测 YOLOv5 – 自定义网络结构（YOLOv5-ShuffleNetV2）
目标检测 YOLOv5 – 常见的边框（bounding box ）坐标表示方法
目标检测 YOLOv5 – 图像大小与loss权重的关系
目标检测 YOLOv5 – 根据配置改变网络的深度和宽度
目标检测 YOLOv5 – 转ncnn移动端部署
目标检测 YOLOv5 – Backbone中的Focus
目标检测 YOLOv5 – 模型训练、推理、导出命令
目标检测 YOLOv5 – 人脸数据集widerface转YOLOv5格式
目标检测 YOLOv5 – 使用的数据集格式
目标检测 YOLOv5 – 中使用COCO数据集
目标检测 YOLOv5 – CrowdHuman数据集格式转YOLOv5格式