【YOLO】YOLOv4训练自己的数据集

虽然在笔者写这篇博客的时候，YOLOv5已经问世，但是看到许多博客都说YOLOv4效果更佳，具体如何，笔者得亲测才能认同。
YOLOv4论文链接
YOLOv4代码链接

1 环境搭建

1. 创建虚拟环境并激活
conda create -n yolo4 pip python=3.6
conda activate yolo4
2. Requirements 安装必要的依赖项
根据github：

 conda install cudatoolkit=10.0
 conda install cudnn=7.6.5
 conda install cmake  # 这里笔者安装的cmake=3.18.2
 conda install opencv  # 这里笔者安装的opencv=3.4.2
1
2
3
4

笔者选择的框架是tensorflow，如果电脑允许最好下载GPU版本的tensorflow

根据github上的指示，进入tensorflow框架下的yolov4源码github的链接：



根据指示继续安装requirements

 pip install --ignore-installed --upgrade tensorflow-gpu==2.3.0 -i https://pypi.douban.com/simple
 pip install opencv-python==4.1.1.26; lxml; tqdm...
1
2

3. 下载源代码

git clone https://github.com/AlexeyAB/darknet.git
1

编译

由于后续训练要使用GPU，因此需要修改makefile文件，具体修改部分如图所示：

根据上图的标注部分进行修改

darknet YOLO 编译使用GPU

• NVCC要对应自己的电脑路径进行修改

要根据电脑的计算能力修改，笔者的电脑无法满足，于是将其注释掉,选择可以满足的计算能力，可以百度搜索自己电脑的GPU 计算能力。

在终端运行，编译：

1 # cd到darknet文件夹下： 
2 make # 或make -j8
1
2

2 下载开源权重并测试

yolov4.weights官网链接
笔者百度网盘链接：https://pan.baidu.com/s/1h4vuNFgWRyiMoppNnVR1wA
提取码：h0c5

./darknet/results detect cfg/yolov4.cfg yolov4.weights data/dog.jpg
1

其中，./darknet/results结果保存路径；data/dog.jpg测试图片

3 训练自己的数据集

数据集结构如下图所示：

其中，Annotations: 标签文件（.xml）
ImageSets/Main: 在三个.txt中存放三类数据集的图片名称
将train_val_test.py代码复制到ImageSets相同根目录下，生成 需要的.txt文件，代码如下：
可以自行修改数据集的比例。

"""
2020.09.18:alian
divide train&val&test save as txt
"""
import os
import random
from pathlib import Path

trainval_percent = 0.8  # trainval数据集占所有数据的比例
train_percent = 0.75  # train数据集占trainval数据的比例
xmlfilepath = 'Annotations'  # 标注数据xml储存的路径
txtsavepath = 'ImageSets/Main'  # txtb保存的路径
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
print('total number is ', num)
list = range(num)
tv = int(num * trainval_percent)
print('trainVal number is ', tv)
tr = int(tv * train_percent)
print('train number is ', tr)
print('test number is ', num - tv)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

# windows
# ftrainval = open(txtsavepath+'\\'+'trainval.txt', 'w')
# ftest = open(txtsavepath+'\\'+'test.txt', 'w')
# ftrain = open(txtsavepath+'\\'+'train.txt', 'w')
# fval = open(txtsavepath+'\\'+'val.txt', 'w')
# ubuntu
ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

JepgImages: 原始图片

1. 准备YOLOv4需要的label和txt
./darknet/scripts下的voc_label.py复制到项目根目录下./darknet，并对其内容进行修改:

"""
2020.11.09
author:alian
"""
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

# sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
sets=[('turnout', 'train'), ('turnout', 'val'),  ('turnout', 'test')]  # 以类别名称命名
# classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
classes = ["turnout"]  # 类别名称

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(classes, image_id):
    # in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
    # out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
    in_file = open('VOC2007_%s/Annotations/%s.xml'%(classes, image_id))
    out_file = open('VOC2007_%s/labels/%s.txt'%(classes, image_id), 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()

for classes, image_set in sets:
    if not os.path.exists('VOC2007_%s/labels/'%(classes)):
        os.makedirs('VOC2007_%s/labels/'%(classes))
    image_ids = open('VOC2007_%s/ImageSets/Main/%s.txt'%(classes, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(classes, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOC2007_%s/JPEGImages/%s.jpg\n'%(wd, classes, image_id))
        convert_annotation(classes, image_id)
    list_file.close()

# os.system("cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt > train.txt")
# os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt")
os.system("cat turnout_train.txt turnout_val.txt > train.txt")
os.system("cat turnout_train.txt turnout_val.txt turnout_test.txt > train.all.txt")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

执行如下命名：
在根目录下将会生成训练需要的文件，即各个训练集中包含图像的路径。

cd ./darknet
python voc_label.py
1
2

2、修改配置文件
（1）. cfg目录下：（voc.data/coco.data二选一即可，笔者这里选择cov.data）

• 类别数量
• 训练过程中训练数量和验证数量的txt文件（即上述voc_label.py生成的文件）
• 类别标签名称
• 存放权重的路径

（2）yolov-obj.cfg

• yolov4训练参数和相关网络结构的修改：
  复制cfg文件夹下的yolov4-custom.cfg重命名为yolo-obj.cfg，修改如下参数：
  输入图像大小和训练测试阶段中的batch的数量和划分次数；
  图像大小必须是32的倍数
  

• 训练代数：

• github中给出了max_batches的基本设置方法，2000 ×classes（但不少于训练图像的总数以及6000）。当然，设置的大一些也是可以的，只不过后期基本上在某一值附近震荡

• 值得注意的是，steps的设置是max_batches × 80% 和 max_batches × 90%。
  

• 网络结构：

根据待测目标类别的数量更改YOLO层(3处地方)和YOLO层前一层的卷积层（3处地方）

• 包含YOLO前一层卷积层的卷积核个数：（classes + 5）*3 （三处）
• YOLO层的类别数classes （三处）
• 锚框（可选，kmeans聚类） （暂时不改）
  

（3） data目录下：

开始训练

下载预训练权重：yolov4.conv.137
笔者百度云盘链接：https://pan.baidu.com/s/1p-7zbIIDNqQrOZGAWG3B6Q
提取码：zoiw
放到./darknet文件夹下面

• 训练指令：

1./darknet detector train cfg/voc.data cfg/yolo-obj.cfg yolov4.conv.137 -gpus 6,7  # 多GPU训练
1

训练好的权重文件将会保存在./darknet/backup文件夹下面

• 若在服务器上训练，无法可视化损失窗口，则使用如下命令：

1./darknet detector train cfg/voc.data cfg/yolo-obj.cfg yolov4.conv.137 -dont_show
1

• 在没有GUI的远程服务器上进行训练期间查看mAP和损失图，使用命令

./darknet detector train cfg/voc.data cfg/yolo-obj.cfg yolov4.conv.137 -dont_show -mjpeg_port 8090 -map
1

在谷歌或者火狐浏览器上查看http://ip-address:8090

• 继续训练指令：

./darknet detector train cfg/voc.data cfg/yolo-obj.cfg backup\yolo-obj_2000.weights
1

其中：backup\yolo-obj_2000.weights为最后的模型

预测

1 ./darknet detector test cfg/voc.data cfg/yolov4-custom.cfg yolov4-custom_xxxx.weights
1