YOLOv5-Shufflenetv2

YOLOv5 General steps to modify the network structure in ：

models/common.py： stay common.py In file , Add the module code to be modified
models/yolo.py： stay yolo.py In file parse_model Add the name of the new module to the function
models/new_model.yaml： stay models Corresponding to the new module under the folder .yaml file

One 、Shufflenetv2
[Cite]Ma, Ningning, et al. “Shufflenet v2: Practical guidelines for efficient cnn architecture design.” Proceedings of the European conference on computer vision (ECCV). 2018.
Open vision lightweight convolutional neural network Shufflenetv2, Four lightweight network design criteria are proposed through a large number of experiments , For input and output channels 、 Number of group convolution groups 、 The degree of network fragmentation 、 The speed and memory access of element by element operation on different hardware MAC(Memory Access Cost) The impact is analyzed in detail ：

Rule one ： The number of input and output channels is the same , Memory accesses MAC Minimum
Mobilenetv2 Not satisfied , The quasi residual structure is adopted , The number of input and output channels is not equal
Guideline 2 ： If the number of packets is too large, the packet convolution will increase MAC
Shufflenetv1 Not satisfied , We use group convolution （GConv）
Rule three ： Fragmentation operation （ Multi channel , Make the network very wide ） Unfriendly to parallel acceleration
Inception Series of networks
Guideline 4 ： Operate element by element （Element-wise, for example ReLU、Shortcut-add etc. ） The resulting memory and time consumption cannot be ignored
Shufflenetv1 Not satisfied , Adopted add operation
For the above four guidelines , The author puts forward Shufflenetv2 Model , adopt Channel Split Replace packet convolution , Meet four design criteria , The best trade-off between speed and accuracy is achieved .
Model overview

Shufflenetv2 There are two structures ：basic unit and unit from spatial down sampling(2×)

basic unit： The number of input and output channels remains unchanged , The size remains the same
unit from spatial down sample ： Double the number of output channels , Double the size （ Downsampling ）
Shufflenetv2 The overall philosophy should closely approach the four principles of lightweight proposed in the paper , In addition to basic rule 4 , Have effectively avoided .

In order to solve GConv(Group Convolution) Resulting in different group There is no information exchange between , Only in the same group The problem of feature extraction in ,Shufflenetv2 Designed Channel Shuffle Operation for channel rearrangement , Span group Information exchange

class ShuffleBlock(nn.Module):
    def __init__(self, groups=2):
        super(ShuffleBlock, self).__init__()
        self.groups = groups

    def forward(self, x):
        '''Channel shuffle: [N,C,H,W] -> [N,g,C/g,H,W] -> [N,C/g,g,H,W] -> [N,C,H,W]'''
        N, C, H, W = x.size()
        g = self.groups
        return x.view(N, g, C//g, H, W).permute(0, 2, 1, 3, 4).reshape(N, C, H, W)

Join in YOLOv5
common.py File modification ： Add the following code directly at the bottom

# ---------------------------- ShuffleBlock start -------------------------------

#  Channel rearrangement , Span group Information exchange 
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)  #  By dimension 1 Conduct 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

# ---------------------------- ShuffleBlock end --------------------------------

yolo.py File modification ： stay yolo.py Of parse_model Function , Join in conv_bn_relu_maxpool, Shuffle_Block Two modules （ As shown in the red box below ）

newly build yaml file ： stay model New under file yolov5-shufflenetv2.yaml file , Just copy the following code

# YOLOv5  by Ultralytics, GPL-3.0 license

# Parameters
nc: 20  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # 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]
  # Shuffle_Block: [out, stride]
  [[ -1, 1, conv_bn_relu_maxpool, [ 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
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P4
   [-1, 1, C3, [256, 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, [256, False]],  # 17 (P4/16-medium)

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

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

reference ：https://blog.csdn.net/weixin_43799388/article/details/123597320