yolov5更换MobileNetV3的主干网络

目录

1.定义网络

 2.将test_MobileNetV3.py上面的代码复制粘贴到如下文件里面

2.1复制需要更改的原来的网络如何改结构

3.更改yolo.py注册网络 

1.定义网络

1. 宽度为0.5的YOLOV5网络的结构图

在主干网络上面可以重新定义成三层，编号从0开始

 如图是MobileNetV3 的网络结构，要想重新定义的画需要保持每次输出图片的大小不变

 定义MobileNetV3 的代码如下，我们可以分为3层 test_MobileNetV3.py

import torch
from torch import nn
from torchvision import models
from torchinfo import summary
class MobileNetV3(nn.Module):
    def __init__(self, n):
        super().__init__()
        model = models.mobilenet_v3_small(pretrained=True)
        if n == 0:
            self.model = model
        if n == 1:
            self.model = model.features[:4]
        if n == 2:
            self.model = model.features[4:9]
        if n == 3:
            self.model = model.features[9:]

    def forward(self, x):
        return self.model(x)

if __name__ == '__main__':
    x = torch.randn(1, 3, 640, 640)
    net = MobileNetV3(0)
    out = net(x)
    print(x.shape)
    summary(net,(1,3,640,640))

feature代表的含义 

 

 2.将test_MobileNetV3.py上面的代码复制粘贴到如下文件里面

2.1复制需要更改的原来的网络如何改结构

从这些里面挑一个

做出如下更改

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 80  # 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
backbone:
  # [from, number, module, args]
  [
    [-1, 1,MobileNetV3, [24, 1]],  # 0-P1/2
   [-1, 1,MobileNetV3, [48,2]],  # 1-P2/4
   [-1, 1,MobileNetV3, [576,3]],
  ]

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

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

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

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

下面几张图是解释

 

3.更改yolo.py注册网络 

如下位置

 添加如下代码，大概340行左右

还是在yolo.py文件里面更改这一句可以进行测试

run一下yolo.py

 更改完成！！！