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 

标签：__,主干,yolov5,shufflent,self,layer,v2,model,x0	

From： https://blog.csdn.net/m0_53291740/article/details/141022127