yolov5 修改

1. 修改网络结构----以yolov8 的C2F为例

借鉴的代码：http://github.com/ultralytics/ultralytics

修改顺序：

1）models/commons.py ------------ 加入新增网络结构

打开models，复制里面的c2f模块，到yolov5的common.py里面。

class C2f(nn.Module):
    # CSP Bottleneck with 2 convolutions
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    def forward(self, x):
        y = list(self.cv1(x).chunk(2, 1))
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

    def forward_split(self, x):
        y = list(self.cv1(x).split((self.c, self.c), 1))
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

 注意这里有个Bottleneck，yolov5里也有用，但是不一样，v8多了k这个参数，所以还要复制yolov8里的Bottleneck

class Bottleneck(nn.Module):
    # Standard bottleneck
    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):  # ch_in, ch_out, shortcut, groups, kernels, expand
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, k[0], 1)
        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

这里为了和yolov5里面的区分，避免其他模块的运用，修改名称，在前面添加c2f，为c2

fBottleneck。

2）models/yolo.py -------------------------------  设定网络结构的传参细节

在yolov5的yolo.py文件的 parse_model里面添加c2f。

3）models/yolov5*.yaml --------------------- 修改现有模型结构配置文件

将C3替换为c2f。可以backbone和head都改了。

4）train.py ----------------------------- 训练时指定模型结构配置文件

yolov5s.pt虽然结构和改过的不一样，但是也可以用，预训练的权重，它会自己去加载和迁移，不能用的权重，它不会全部迁移。 

 2.引入注意力机制----以SE为例

借鉴的代码：https://github.com/ZhugeKongan/Attention-mechanism-implementation

1）models/commons.py ------------ 加入新增网络结构

SE_block.py里面

import torch.nn as nn
import torch.nn.functional as F


class SE(nn.Module):

    def __init__(self, in_chnls, ratio):
        super(SE, self).__init__()
        self.squeeze = nn.AdaptiveAvgPool2d((1, 1))
        self.compress = nn.Conv2d(in_chnls, in_chnls // ratio, 1, 1, 0)
        self.excitation = nn.Conv2d(in_chnls // ratio, in_chnls, 1, 1, 0)

    def forward(self, x):
        out = self.squeeze(x)
        out = self.compress(out)
        out = F.relu(out)
        out = self.excitation(out)
        return x*F.sigmoid(out)

全部copy进自己的common.py里面。

2）models/yolov5*.yaml --------------------- 修改现有模型结构配置文件

 注意：当引入新的层时，要修改后续的结构中的from参数。

加到最后， SE的输入和输出是一样的，1024，再加一个隐层的输出，ratio。

要注意添加了一个层，后面层的序号就变了，相应的contact所连接的层也变了。

例如cat head P4，原来是第14层，改为第15层。

cat head P5，原来是第10层，改为第11层。

3）models/yolo.py --------------- 设定网络结构中的传参细节

注意：当新的自定义模块中存在输入输出维度时，要使用gw调整输出维度

        elif m is SE:
            c1 = ch[f]
            c2 = args[0]
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, 8)
            args = [c1, *args[1:]]

 

 4）train.py ----------------------------- 训练时指定模型结构配置文件

 3.替换主干网络----以MobileNet为例

借鉴的代码：torchvision/ timm

 繁杂，不好看，安装一个库

pip install torchinfo

 导入包

 做迁移时，只用到它特征提取的部分，即它的第一个sequential。

添加 features

1）models/commons.py -------------------- 加入新增网络结构

Sequential可以通过切片来获取模型结构

导包并添加模块

import torchvision.models as models

class MobileNetV3(nn.Module):

    def __init__(self, slice):
        super(MobileNetV3, self).__init__()
        self.model = None
        if slice == 1:
            self.model = models.mobilenet_v3_small(pretrained=True).features[:4]
        elif slice == 2:
            self.model = models.mobilenet_v3_small(pretrained=True).features[4:9]
        else:
            self.model = models.mobilenet_v3_small(pretrained=True).features[9:]

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

 2）models/yolov5*.yaml ------------------- 修改现有模型结构配置文件

backbone:
  # [from, number, module, args]
  [
  # [-1, 1, Focus, [64, 3]],  # 0-P1/2
  # [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
  #  [-1, 3, C3, [128]],
  #  [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
  #  [-1, 9, C3, [256]],
  #  [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
  #  [-1, 9, C3, [512]],
  #  [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
  #  [-1, 1, SPP, [1024, [5, 9, 13]]],
  #  [-1, 3, C3, [1024, False]],  # 9
    [-1, 1, MobileNetV3, [24,1]],
    [-1, 1, MobileNetV3, [48,2]],
    [-1, 1, MobileNetV3, [576,3]]
  ]

# YOLOv5 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）models/yolo.py ------------------ 设定网络结构中的传参细节

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

 4）train.py ---------------------------- 训练时指定模型结构配置文件