1、准备工作
首先准备训练用到的文件,放在自己的谷歌云盘
1、自己的数据集(建议使用压缩包上传,上传的速度和加载速度会快很多)yolo训练使用的是txt格式的标签,里面是目标的坐标值。若已标的标签是xml文件,可以使用此Python代码提取对应的坐标信息生成TXT文件。
2、data文件,names文件
data文件放置类别数和一些路径,names文件放置类别名
我的是obj.data和obj.names
注:这里的backup文件夹放置训练后的权重文件,本来应该将目录放在我的硬盘下以防自动断开文件丢失,但每次生成权重文件时就会报找不到目录,所以只能放在其他目录,为防止自动断开,还是要时常复制一下
3、yolov4.weights(测试darknet使用)和yolov4.conv.137(首次训练使用)
2、装载yolov4
详细解释可参考另一篇博客(使用Google Colab装载YOLOv4(免费云端GPU,无需配置任何环境))
命令如下:
1、查看GPU
! nvidia-smi
2、从GitHub下载darknet(colab配合GitHub使用,速度非常快)
! git clone https://github.com/AlexeyAB/darknet.git
3、修改Makefile相关参数
! sed -i 's/GPU=0/GPU=1/g' /content/darknet/Makefile
! sed -i 's/CUDNN=0/CUDNN=1/g' /content/darknet/Makefile
! sed -i 's/OPENCV=0/OPENCV=1/g' /content/darknet/Makefile
4、编译darknet
! cd darknet; make
5、将权重拷贝到content目录(工作的默认根目录)
!cp /content/drive/My\ Drive/xxx/yolov4.weights ./
6、测试是否装载成功
# 检测试一下
%cd /content/darknet
!./darknet detect cfg/yolov4.cfg /content/yolov4.weights data/dog.jpg
# 上述函数写好了,显示一下
from PIL import Image
Image.open('/content/darknet/predictions.jpg')
3、训练自己的数据集
1、复制相关文件到工作区
首先自己在content目录下新建data文件夹,在data文件夹下面新建backup文件夹,如图
执行命令
%cd /content
!cp /content/drive/My\ Drive/xxx/yolov4-custom.cfg ./yolov4-custom.cfg
!cp /content/drive/My\ Drive/xxx/yolov4.conv.137 ./
!cp /content/drive/My\ Drive/xxx/obj.data ./data/
!cp /content/drive/My\ Drive/xxx/obj.names ./data/
2、复制数据集并解压缩
!cp /content/drive/My\ Drive/robomaster/train01/armor01.zip ./
!unzip ./armor01.zip -d data
3、自动生成train.txt文件
%cd /content
import os
image_files = []
os.chdir(os.path.join("data", "armor01"))
for filename in os.listdir(os.getcwd()):
if filename.endswith(".jpg"):
image_files.append("/content/data/armor01/" + filename)
os.chdir("..")
with open("train.txt", "w") as outfile:
for image in image_files:
outfile.write(image)
outfile.write("\n")
outfile.close()
os.chdir("..")
若还需要验证和测试文件,仿照此代码添加即可
4、开始训练
# train your custom detector
!./darknet/darknet detector train /content/data/obj.data /content/yolov4-custom.cfg /content/yolov4.conv.137 -dont_show
若训练中途断掉需要在次训练,可以使用之前训练的权重在次训练。
!./darknet/darknet detector train /content/data/obj.data /content/yolov4-custom.cfg /content/drive/My\ Drive/robomaster/train01/backup/yolov4-custom_final.weights -dont_show
5、训练好了可以测试一下
首先将cfg文件中的训练值改过来
# 设置配置文件
%cd /content
!sed -i 's/batch=64/batch=1/' yolov4-custom.cfg
!sed -i 's/subdivisions=64/subdivisions=1/' yolov4-custom.cfg
测试(从B站下的视频)
!./darknet/darknet detector demo -out_filename output.mp4 /content/data/obj.data /content/yolov4-custom.cfg /content/drive/My\ Drive/robomaster/train01/backup/yolov4-custom_last.weights /content/drive/My\ Drive/xxx/20200704212932.mp4 -thresh 0.2 -dont_show