使用C版YOLOv4在自己的数据集上训练测试
本文记录了如何在Ubuntu/Docker中使用Alexey实现的C版YOLOv4在自己的数据集上进行训练与测试。
论文 :
YOLOv4: Optimal Speed and Accuracy of Object Detection
代码 :
https://github.com/AlexeyAB/darknet
环境配置
建议使用docker容器配置环境,docker的安装不在此赘述,docker安装好后拉取一个nvidia/cuda镜像,docker的一些用法可以看这里。如果你不使用docker的话可以跳过docker相关部分。
首先把我们的docker跑起来
$docker run -itdp 0.0.0.0:32774:22 --gpus '"device=1,2"' --name yolo -v /home/usr/:/home -v /data0:/data0 nvidia/cuda-10.0-cudnn7-devel-ubuntu16.04:v1 /bin/bash
"-itdp 0.0.0.0:32774:22"
表示在后台运行一个交互式容器并将主机的32774
端口映射到docker容器的22/tcp
端口,如果不需要指定端口可以将其改为"-itdP"
,这表示我们不指定tcp端口映射而是随机映射。
"--gpus '"device=1,2"
表示我们使用主机的1、2号GPU。docker的gpu命令用法可以看这里。
"--name yolo"
表示容器名称,可自定义。
"-v /home/usr/:/home -v /data0:/data0"
表示将物理机的/home/usr/
目录挂载在docker容器的/home
路径下,后面同理,需要挂载几个目录就写几个-v
参数。
"nvidia/cuda-10.0-cudnn7-devel-ubuntu16.04:v1"
表示要运行的容器名称:容器标签
。
"/bin/bash"
表示容器内使用的命令程序。
docker的一些常用命令用法可以看这里。
接着我们安装依赖
CUDA : 需要CUDA 10.0,安装看这里。
cuDNN : 需要和CUDA 10.0对应的cuDNN,cuDNN在这里获取,选择最新的cuDNN for CUDA 10.0并下载cuDNN Library for Linux
,安装看这里。
CMake : 需要CMake>=3.8,CMake在这里获取,选择"Unix/Linux Source (has \n line feeds)"
后对应的文件下载,使用如下命令编译安装
$tar xvzf cmake-3.17.1.tar.gz # 注意替换你下载的版本号,下同
$cd cmake-3.17.1
$./configure
$make -j
$make install
OpenCV : 需要OpenCV>=2.4,OpenCV在这里获取,选择Sources
下载,使用如下命令编译安装
$unzip opencv-4.2.0.zip # 注意替换成你下载的文件名,下同
$cd opencv-4.2.0.zip
$mkdir build
$cd build
$cmake ..
$make -j
$make install
上述各项依赖安装期间有权限问题则使用sudo
;各项安装完成后应查看版本号;必要时通过修改~/.bashrc
添加相关的环境变量并使用source ~/.bashrc
激活,此处不再赘述。
然后我们编译YOLOv4的darknet
将项目克隆至本地
$git clone https://github.com/AlexeyAB/darknet
$cd darknet
使用下列方式之一进行编译
- 使用CMake
$mkdir build-release
$cd build-release
$cmake ..
$make -j
$make install
如果期间报错,可能是之前的依赖安装有问题,按错误提示Google
解决,然后再次编译,注意再次编译前使用make clean
清除之前编译失败后遗留的中间文件。
- 使用make
在make
前可修改darknet/Makefile
文件前几行的相关参数,参数含义在这里,一般至少将GPU
和CUDNN
设置为1,以使用GPU加速。
$make -j
如果有报错,解决后报错后同样记得make clean
。
编译完成后执行
$./darknet
usage: ./darknet <function>
有如上结果说明编译成功。
最后我们简单测试一下编译好的darknet
将yolov4.weights下载到darknet
目录中,可能需要机智上网。
Google Drive 地址 :
https://drive.google.com/open?id=1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT
GitHub 地址 :
https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights
然后执行
$./darknet detect cfg/yolov4.cfg yolov4.weights data/dog.jpg
CUDA-version: 10000 (10010), cuDNN: 7.6.5, GPU count: 2
OpenCV version: 4.2.0
compute_capability = 750, cudnn_half = 0
net.optimized_memory = 0
mini_batch = 1, batch = 8, time_steps = 1, train = 0
...
...
[yolo] params: iou loss: ciou (4), iou_norm: 0.07, cls_norm: 1.00, scale_x_y: 1.05
nms_kind: greedynms (1), beta = 0.600000
Total BFLOPS 128.459
avg_outputs = 1068395
Allocate additional workspace_size = 52.43 MB
Loading weights from yolov4.weights...
seen 64, trained: 32032 K-images (500 Kilo-batches_64)
Done! Loaded 162 layers from weights-file
data/dog.jpg: Predicted in 103.595000 milli-seconds.
bicycle: 92%
dog: 98%
truck: 92%
pottedplant: 33%
...
可以看到最后输出的检测结果等,目录下会生成预测结果。
训练
按下面的格式组织我们的数据
最后用于训练的命令是
./darknet detector train xxx.data xxx.cfg xxx.weights
因此在训练前我们需要先组织好三个文件: 用于描述数据集信息的.data文件、用于描述网络信息的.cfg文件和一个.weights预训练权重文件。另外我们需要为数据集准备labels文件。
- 下载预训练权重
在下面的地址之一下载预训练权重,将其放入darknet目录。
-
Google Drive地址
https://drive.google.com/open?id=1JKF-bdIklxOOVy-2Cr5qdvjgGpmGfcbp -
GitHub地址
https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137
- 调整.data文件,文件内容示例如下
classes= 2
train = data/train.txt
valid = data/test.txt
names = data/obj.names
backup = backup/
- 将
classes
设置为你数据集的目标类别数。 - 将
train
设置为训练集文件train.txt
的路径,train.txt
中包含所有的训练集数据路径,每行一条,注意相对路径或绝对路径均可,相对路径是相对于可执行文件darknet
的,如下所示:
data/obj/img1.jpg
data/obj/img2.jpg
data/obj/img3.jpg
- 将
valid
设置为测试集文件test.txt
的路径,方法同train
。 - 将
names
设置为一个.names
文件的路径,.names
文件中包含了你数据集目标的所有类别,每行一个,如下所示:
dog
cat
person
...
- 将
backup
设置为你将要存放中间权重的目录路径,训练过程中每100个iteration会更新一次latest权重,每10000个iteration会保存一次中间权重,这些权重文件将存储在你设置的backup
目录中。
- 调整.cfg文件。
我们首先将cfg/yolov4-custom.cfg
复制一份并重命名为yolov4-obj.cfg
(或其他名字),然后基于yolov4-obj.cfg
进行修改。
yolov4-obj.cfg
文件中,开头部分描述了用于训练网络的一些超参数,然后描述了完整的网络结构。我们将修改一部分超参数和网络结构使其适应我们的数据集和训练环境。
- 设置
batch=64
。 - 设置
subdivisions=16
,如果稍后训练时显存溢出的话,可将此处调整为32、64。 - 设置
max_batches=classes*2000
,但是不要小于训练集的数据量,也不要小于6000,比如你有5类目标,则设置max_batches=10000
。 - 设置
steps
为max_batches
的80%, 90%。 - 设置
width=416
,height=416
或任何32的整数倍,根据数据特点和显存容量决定。 - 将每个
[yolo]
层的classes
设置为你的数据集目标类别数。注意一共有3个[yolo]
层。 - 将每个
[yolo]
层的上一层中的filters
设置为(classes + 5)x3
,其中classes
就是刚才设置的类别数。
- 准备label文件。
我们需要为每个数据图像准备一份label文件,里面每一行描述该图像中标注的一个object的类别和bbox信息,如下所示:
<object-class> <x_center> <y_center> <width> <height>
其中object-class
# TODO: Detail
- 最后,我们的数据集应该组织为如下形式:
Annotations
JPEGImages
labels
trainval.txt
test.txt
其中Annotations
中放置所有图片的.xml标注文件,JPEGImages
中放置所有的图片,labels
中放置上文所述的label文件,trainval.txt
和test.txt
为上文所述训练集和测试集文件。
开始训练
# TODO: Start Train
测试
# TODO: Test