原始的yolov5s的网络结构如下所示:
1 修改任务
现对模型结构进行修改,将网络结构修改为如下图所示:
其中C2是在C3的基础上进行的修改,其结构如下图所示:
经分析可知,需要做如下修改:
(1)将第4层C3*2修改为C2*2
(2)将第6层C3*3修改为C3*1
(3)去除第7、8层
2 模型配置文件(yaml)修改
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args] args:channel、kernel_size、stride、padding、bias
[[-1, 1, Conv, [64, 6, 2, 2]], # 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, 6, C2, [128]], # 将C3*2修改为C2*2
# [-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 3, C3, [512]], # 将C3*3修改为C3*1
# [-1, 9, C3, [512]],
# [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 # 删除第7层
# [-1, 3, C3, [1024]], # 删除第8层
[-1, 1, SPPF, [512, 5]], # 9 # 为连接第6层,修改相应参数
# [-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
# [[-1, 1, Conv, [512, 1, 1]],
[[-1, 1, Conv, [512, 3, 2]],# 修改参数
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 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, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
# [[-1, 14], 1, Concat, [1]], # cat head P4
[[-1, 12], 1, Concat, [1]], # cat head P4 # 修改层数 14-2
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
# [[-1, 10], 1, Concat, [1]], # cat head P5
[[-1, 8], 1, Concat, [1]], # cat head P5 # 修改层数 10-2
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
# [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
[[15, 18, 21], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) # 修改层数
]
根据分析,将yolov5s.yaml文件进行如上所示修改,注意由于删除了两个层,在Head中,部分输入层需要进行相应的修改。
另外,C2的实现已在之前的文章中有作详细说明。
3 训练
经过修改后,执行以下命令进行训练,
train.py --img 640 --batch 4 --epoch 1 --data data/fruit.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt
训练时的部分信息如下所示:
以上仅为测试,查看之前的修改是否正确,从训练结果看,mAP值很低,需要增加训练。