yolov5更换主干网络shufflent

目录

1.网络结构解析

1.1创建yolov5s_shufflent_v2_X0_5.yaml文件

2.对common.py末尾进行添加

 3.修改yolo.py

---

1.网络结构解析

1.可以先看看shufflenet_v2的网络结构

import torch
from torch import nn
from torchvision import models
from torchinfo import summary


class shufflenet_v2_x0_5(nn.Module):
    def __init__(self,n):
        super().__init__()
        model = models.shufflenet_v2_x0_5(pretrained=True)
        self.model=model
    def forward(self, x):
        return self.model(x)
if __name__ == '__main__':
    x=torch.randn(1,3,640,640)
    net=shufflenet_v2_x0_5(0)
    out=net(x)
    print(out.shape)
    summary(net,(1,3,640,640))

这个是YOLOV5的网络。框出来的是yolov5的主干网络。我们用shufflenet_v2的部分替换。可以直接把shufflenet_v2的网络截取出三部分

定义

 下图的右边部分是网络shufflenet的官方网络结构，直接使用即可。

定义我们自己需要修改的shufflenet类

import torch
from torch import nn
from torchvision import models
from torchinfo import summary


class Shufflenet_v2_x0_5(nn.Module):
    def __init__(self,n):
        super().__init__()
        model = models.shufflenet_v2_x0_5(pretrained=True)
        if n==1:
            layer=[]
            layer+=[model.conv1]
            layer+=[model.maxpool]
            layer+=[model.stage2]
            self.model=nn.Sequential(*layer)
        if n==2:
            self.model=model.stage3
        if n==3:
            layer=[]
            layer+=[model.stage4]
            layer+=[model.conv5]
            self.model = nn.Sequential(*layer)
    def forward(self, x):
        return self.model(x)
if __name__ == '__main__':
    x=torch.randn(1,3,640,640)#torch.Size([1, 48, 80, 80])
    net=Shufflenet_v2_x0_5(1)
    out=net(x)
    print(out.shape)

    x1=torch.randn(1,48,80,80)#torch.Size([1, 96, 40, 40])
    net1 = Shufflenet_v2_x0_5(2)
    out1 = net1(x1)
    print(out1.shape)
    x2=torch.randn(1, 96, 40, 40)#torch.Size([1, 1024, 20, 20]
    net2 = Shufflenet_v2_x0_5(3)
    out2 = net2(x2)
    print(out2.shape)
    # summary(net,(1,3,640,640))

1.1创建yolov5s_shufflent_v2_X0_5.yaml文件

照着上面的网络对齐修改

# 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,Shufflenet_v2_x0_5, [48, 1]],  # 0-P1/2
   [-1, 1,Shufflenet_v2_x0_5, [96,2]],  # 1-P2/4
   [-1, 1,Shufflenet_v2_x0_5, [1024,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)
  ]

2.对common.py末尾进行添加

位置如下：

        

class Shufflenet_v2_x0_5(nn.Module):
    def __init__(self,n):
        super().__init__()
        model = models.shufflenet_v2_x0_5(pretrained=True)
        if n==1:
            layer=[]
            layer+=[model.conv1]
            layer+=[model.maxpool]
            layer+=[model.stage2]
            self.model=nn.Sequential(*layer)
        if n==2:
            self.model=model.stage3
        if n==3:
            layer=[]
            layer+=[model.stage4]
            layer+=[model.conv5]
            self.model = nn.Sequential(*layer)
    def forward(self, x):
        return self.model(x)

 3.修改yolo.py

解析参数，并运行

 elif m is Shufflenet_v2_x0_5:
            c2=args[0]
            args=args[1:]

'yolov5s_shufflent_v2_X0_5.yaml'

打印网络参数如下：