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:]