【人工智能学习之YOLOV5更换主干网络】

GhostNet

更换GhostNet主干网络比较简单，因为models/common.py里面已经写好了相关类的定义，如下图

所以我们直接配置yaml文件即可

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license

# Parameters
nc: 7 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # 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
# ghostnet
backbone:
  # [from, number, module, args]
  [
    [-1, 1, GhostConv, [64, 3, 2]], # 0-P1/2
    [-1, 1, GhostConv, [128, 3, 2]], # 1-P2/4
    [-1, 3, C3Ghost, [128]],
    [-1, 1, GhostConv, [256, 3, 2]], # 3-P3/8
    [-1, 6, C3Ghost, [256, 1]],
    [-1, 1, GhostConv, [512, 3, 2]], # 5-P4/16
    [-1, 9, C3Ghost, [512, 1]],
    [-1, 1, GhostConv, [1024, 3, 2]], # 7-P5/32
    [-1, 3, C3Ghost, [1024, 1]],
    [-1, 1, SPPF, [1024, 5]], # 9
  ]

# YOLOv5 v6.0 head
head: [
    [-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 6], 1, Concat, [1]], # cat backbone P4
    [-1, 3, C3, [512, False]], # 13

    [-1, 1, Conv, [256, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 4], 1, Concat, [1]], # cat backbone P3
    [-1, 3, C3, [256, False]], # 17 (P3/8-small)

    [-1, 1, Conv, [256, 3, 2]],
    [[-1, 14], 1, Concat, [1]], # cat head P4
    [-1, 3, C3, [512, False]], # 20 (P4/16-medium)

    [-1, 1, Conv, [512, 3, 2]],
    [[-1, 10], 1, Concat, [1]], # cat head P5
    [-1, 3, C3, [1024, False]], # 23 (P5/32-large)

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

并在yolo.py中进行一点修改：

最后在调用yaml文件时将原本的yoloXX.yaml切换为ghostnet.yaml即可。

    parser.add_argument("--cfg", type=str, default="models/ghostnet.yaml", help="model.yaml path")

ShuffleNet

ShuffleNet的修改较为复杂，需要在models/common.py里面添加相关类的定义，我们直接写在common的末尾就可以：

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 stem(nn.Module):
    def __init__(self, c1, c2):
        super(stem, 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

添加完成之后，我们再进行yaml文件的配置：

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license

# Parameters
nc: 7 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # 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
# shufflenet
backbone:
  # [from, number, module, args]
  [
    [-1, 1, stem, [64]], # 0-P1/2
    [-1, 1, Shuffle_Block, [128, 2]], # 1-P2/4
    [-1, 3, Shuffle_Block, [128, 1]],
    [-1, 1, Shuffle_Block, [256, 2]], # 3-P3/8
    [-1, 6, Shuffle_Block, [256, 1]],
    [-1, 1, Shuffle_Block, [512, 2]], # 5-P4/16
    [-1, 9, Shuffle_Block, [512, 1]],
    [-1, 1, Shuffle_Block, [1024, 2]], # 7-P5/32
    [-1, 3, Shuffle_Block, [1024, 1]],
    [-1, 1, SPPF, [1024, 5]], # 9
  ]

# YOLOv5 v6.0 head
head: [
    [-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 6], 1, Concat, [1]], # cat backbone P4
    [-1, 3, C3, [512, False]], # 13

    [-1, 1, Conv, [256, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 4], 1, Concat, [1]], # cat backbone P3
    [-1, 3, C3, [256, False]], # 17 (P3/8-small)

    [-1, 1, Conv, [256, 3, 2]],
    [[-1, 14], 1, Concat, [1]], # cat head P4
    [-1, 3, C3, [512, False]], # 20 (P4/16-medium)

    [-1, 1, Conv, [512, 3, 2]],
    [[-1, 10], 1, Concat, [1]], # cat head P5
    [-1, 3, C3, [1024, False]], # 23 (P5/32-large)

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

然后在yolo.py中添加：

最后在调用yaml文件时将原本的yoloXX.yaml切换为shufflenet.yaml即可。

    parser.add_argument("--cfg", type=str, default="models/shufflenet.yaml", help="model.yaml path")

MobileNet

MobileNet的添加我们采取另一种方式。
同样还是需要在models/common.py的末尾添加相关的类：

from torchvision import models
class MobileNetV3(nn.Module):
    def __init__(self, n):
        super().__init__()
        model = models.mobilenet_v3_small(pretrained=True)
        if n == 1:
            self.model = model.features[:4]
        elif n == 2:
            self.model = model.features[4:9]
        elif n == 3:
            self.model = model.features[9:]
        else:
            print("其他")
    def forward(self, x):
        return self.model(x)

然后我们再配置yaml和修改yolo.py：

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license

# Parameters
nc: 7  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # 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
# mobilenet
backbone:
  # [from, number, module, args]
  [
    [-1, 1, MobileNetV3, [24, 1]],  # 80 * 80
    [-1, 1, MobileNetV3, [48, 2]],  # 40 * 40
    [-1, 1, MobileNetV3, [576, 3]]  # 20 * 20
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],      # 40 * 40
   [[-1, 1], 1, Concat, [1]],  # cat backbone P4    # 40 * 40
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],      # 80 * 80
   [[-1, 0], 1, Concat, [1]],  # cat backbone P3    # 80 * 80
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],                      # 40 * 40
   [[-1, 7], 1, Concat, [1]],  # cat head P4       # 40 * 40
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

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

   [[10, 13, 16], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

yolo的添加位置有所不同，需要添加俩个地方：

最后在调用yaml文件时将原本的yoloXX.yaml切换为mobilenet.yaml即可。

    parser.add_argument("--cfg", type=str, default="models/mobilenet.yaml", help="model.yaml path")