本文为365天深度学习训练营 中的学习记录博客
原作者:K同学啊|接辅导、项目定制
本次训练是在前文《YOLOv5白皮书-第Y2周:训练自己的数据集》的基础上进行的,并参考了《YOLOv5白皮书-第Y5周:yolo.py文件解读》。
任务:修改了YOLOv5s的网络结构图,请根据网络结构图以及第Y1~Y5周的内容修改对应代码,并跑通程序。
原YOLOv5s的网络结构图:
修改后的YOLOv5s的网络结构图:
通过比较上面两张图,可知,本次的任务:
把索引为4的层从C3 * 2修改为C2 * 2
把索引为6的层从C3 * 3修改为C3 * 1
去除索引为7、8的层
其实还有隐藏任务,就是去除了索引为7、8的层后,原来8层后面的索引也会改变的,这也是要注意的,特别是有concat的地方,更要注意了,这会在后面的代码修改中有解释。
一、代码修改
1、把索引为4的层从C32修改为C22
参考前文《YOLOv5白皮书-第Y5周:yolo.py文件解读》,在common.py中根据下图写出C2的代码。
就是在C3模块代码的基础上修改为C2模块,把C3模块的代码去掉concat后面的Conv,就变成了C2模块。
详见下面的代码,这里把C3模块的代码也写出来,方便和C2模块的代码比较:
python
class C3(nn.Module):
# CSP Bottleneck with 3 convolutions
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
"""Initializes C3 module with options for channel count, bottleneck repetition, shortcut usage, group
convolutions, and expansion.
"""
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
def forward(self, x):
"""Performs forward propagation using concatenated outputs from two convolutions and a Bottleneck sequence."""
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
# 这是在C3模块的基础上修改的,就是把C3模块的self.cv3去掉,主要是把forward(self, x)中的self.cv3去掉
class C2(nn.Module):
# CSP Bottleneck with 2 convolutions
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
"""Initializes C2 module with options for channel count, bottleneck repetition, shortcut usage, group
convolutions, and expansion.
"""
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
# self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
def forward(self, x):
"""Performs forward propagation using concatenated outputs from two convolutions and a Bottleneck sequence."""
return torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1)
在yolo.py中,找到parse_model函数,并在这个函数里的两个地方添加C2,详见注释,如下所示:
python
if m in {
Conv,
GhostConv,
Bottleneck,
GhostBottleneck,
SPP,
SPPF,
DWConv,
MixConv2d,
Focus,
CrossConv,
BottleneckCSP,
C3,
C2, #原代码中没有C2,现在添加了C2
C3TR,
C3SPP,
C3Ghost,
nn.ConvTranspose2d,
DWConvTranspose2d,
C3x,
}:
c1, c2 = ch[f], args[0]
if c2 != no: # if not output
c2 = make_divisible(c2 * gw, ch_mul)
# 下面的代码中原来是没有C2,现在添加了C2(不是指小写的c2)
args = [c1, c2, *args[1:]]
if m in {BottleneckCSP, C3, C2, C3TR, C3Ghost, C3x}: #添加C2
args.insert(2, n) # number of repeats
n = 1
在yolo.py的开头,还要把C2添加上去,详见注释,如下所示:
python
from models.common import (
C3,
C2, #原来没有C2,现在添加C2
C3SPP,
C3TR,
SPP,
SPPF,
Bottleneck,
BottleneckCSP,
C3Ghost,
C3x,
Classify,
Concat,
Contract,
Conv,
CrossConv,
DetectMultiBackend,
DWConv,
DWConvTranspose2d,
Expand,
Focus,
GhostBottleneck,
GhostConv,
Proto,
)
然后在yolov5s.yaml中,在backbone把索引为4的层从C3 * 2修改为C2 * 2 ,因为yolov5s.yaml已经写明depth_multiple为0.33,0.33 * 6 约等于 2,所以第4层为:[-1, 6, C2, [256]]
如下所示:
python
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[
[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]], #2
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C2, [256]], #4 原来是C3*2,现在修改为C2*2
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 3, C3, [512]], #6 原来是C3*3,现在修改为C3*1
#[ -1, 1, Conv, [ 1024, 3, 2 ] ], # 7-P5/32,要注释掉或者删除
#[ -1, 3, C3, [ 1024 ] ], #8,要注释掉或者删除
[-1, 1, SPPF, [1024, 5]], # 9 现在是第7层 ,此层之后的层索引都发生了改变
]
2、把索引为6的层从C3 * 3修改为C3 * 1
因为yolov5s.yaml已经写明depth_multiple为0.33,0.33 * 3 约等于 1,所以第6层为: [-1, 3, C3, [512]],详见下面的代码:
python
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[
[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]], #2
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C2, [256]], #4 原来是C3*2,现在修改为C2*2
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 3, C3, [512]], #6 原来是C3*3,现在修改为C3*1
#[ -1, 1, Conv, [ 1024, 3, 2 ] ], # 7-P5/32,要注释掉或者删除
#[ -1, 3, C3, [ 1024 ] ], #8,要注释掉或者删除
[-1, 1, SPPF, [1024, 5]], # 9 现在是第7层 ,此层之后的层索引都发生了改变
]
3、去除索引为7、8的层
在yolov5s.yaml中,在backbone中把索引为7、8的层注释就可以了,详见下面的代码:
python
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[
[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]], #2
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C2, [256]], #4 原来是C3*2,现在修改为C2*2
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 3, C3, [512]], #6 原来是C3*3,现在修改为C3*1
#[ -1, 1, Conv, [ 1024, 3, 2 ] ], # 7-P5/32,要注释掉或者删除
#[ -1, 3, C3, [ 1024 ] ], #8,要注释掉或者删除
[-1, 1, SPPF, [1024, 5]], # 9 现在是第7层 ,此层之后的层索引都发生了改变
]
但把第7、8层注释掉后,后面的层的索引就变了,在yolov5s.yaml中,head的索引也受到了影响,特别是Concat中,这是用层的索引的,这就是隐藏的任务,所以head的代码修改如下,并参考注释内容:
python
# YOLOv5 v6.0 head
head: [
[-1, 1, Conv, [512, 3, 2]], #8,原10 现在是[512, 3, 2],原来是[512, 1, 1]
[-1, 1, nn.Upsample, [None, 2, "nearest"]], #9,原11
[[-1, 6], 1, Concat, [1]], # cat backbone P4 10,原12, 原12层是cat第6层的,变为10层后无变化
[-1, 3, C3, [512, False]], # 现在是11,原来是13
[-1, 1, Conv, [256, 1, 1]], #12,原14
[-1, 1, nn.Upsample, [None, 2, "nearest"]], #13,原15
[[-1, 4], 1, Concat, [1]], # cat backbone P3 14,原16 原16层是cat第4层的,变为14层后无变化
[-1, 3, C3, [256, False]], # 17 (P3/8-small) 15,原17
[-1, 1, Conv, [256, 3, 2]], #16,原18
[[-1, 12], 1, Concat, [1]], # cat head P4 17,原19 原19层是cat第14层的,但原14层变为第12层,所以要把14修改为12
[-1, 3, C3, [512, False]], # 20 (P4/16-medium) 18,原20
[-1, 1, Conv, [512, 3, 2]], #19,原21
[[-1, 8], 1, Concat, [1]], # cat head P5 #20,原22,原22层是cat原第10层,但原第10层已经变为第8层了,所以要把10修改为8
[-1, 3, C3, [1024, False]], # 23 (P5/32-large) #21,原23
[[15, 18, 21], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 22,原24,原来是[17, 20, 23],层索引已经有变化了,要改为[15, 18, 21]
]
4、补充
因为训练数据的类别是["banana", "snake fruit", "dragon fruit", "pineapple"],只有4种,所以在yolov5s.yaml中还要把nc修改为4,如下所示:
python
# Parameters
nc: 4 # number of classes,原文是80,因为要训练的数据集只有4种类别,所以把80修改为4
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
二、代码运行
开始用自己的数据集训练模型,在项目目录中打开cmd。
如果电脑有GPU,则在cmd中输入命令:
python train.py--img 900 --batch 2 --epoch 50 --data data/ab.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device '0'
如果电脑没有GPU,则在cmd中输入命令:
python train.py --img 900 --batch 2 --epoch 50 --data data/ab.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device cpu
就可以直接训练自己的数据集啦,训练结果如下所示:
三、总结
本次训练需要修改的内容比较多,看起来有三个点,其实是四个点,在yolov5s.yaml中把backbone的索引7、8的层删除了,还要修改head的Concat的索引。