The network structure of the original yolov5s is as follows:
1 Modify task
Now modify the model structure and modify the network structure as shown in the following figure:
Among them, C2 is a modification based on C3, and its structure is shown in the figure below:
After analysis, it can be seen that the following modifications need to be made:
(1) Change layer 4 C3*2 to C2*2
(2) Change layer 6 C3*3 to C3*1
(3) Remove the 7th and 8th layers
2 Model configuration file (yaml) modification
# 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) # 修改层数
]According to the analysis, modify the yolov5s.yaml file as shown above. Note that because two layers are deleted, in the Head, some input layers need to be modified accordingly.
In addition, the implementation of C2 has been described in detail in previous articles.
3 training
After modification, execute the following command for training,
train.py --img 640 --batch 4 --epoch 1 --data data/fruit.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt
Part of the information during training is as follows:
The above is just a test to check whether the previous modification is correct. From the training results, the mAP value is very low, and more training is needed.