YOLOv4+tensorflow2.0训练自己的数据

YOLOv4+tensorflow2.0训练自己的数据

一、环境配置

1.1 实验源码

在YOLOv4原文给出了基于Darknet的实验源码【下载地址】（若链接不可用，请自行复制下面这个网址：https://github.com/AlexeyAB/darknet），该源码核心部分是利用**C++**写的，但是大部分学习目标检测的人都习惯使用python语言，因为python语言提供了大量的框架，更加的简便。因此，要读懂这份源码可能会有一些困难。
为了方便其他研究者，作者在上述源码下方给出了基于其他框架的代码链接，当然也包括了基于tensorflow实现的源码【下载地址】，本篇文章使用的是tensorflow框架中的第一份源码。
在这里插入图片描述
下载后直接解压就行，解压出来是下面这样的：
在这里插入图片描述

1.2 环境要求

在本文使用的源码中，作者明确指出了，改源码使用的是tensorflow2.0版本，所以在实验之前，需要保证自己电脑上装的是2.0版本的tensorflow，如果你的电脑上装的是1.0版本的，则需要提升tensorflow的版本，因为2.0版本和1.0版本在一些函数的调用上会有所不同，如果使用1.0版本的tensorflow来运行该代码，改错会改到让你怀疑人生。所以，还是乖乖的换2.0版本吧！！！
这里我使用的tensorflow版本是2.3.0，搭配的Keras版本是2.3.1
我使用的是anaconda来搭建环境，主要步骤如下：

1.2.1 安装anaconda

anaconda的安装步骤见另一篇文章：Anaconda 安装教程

1.2.2 在anaconda中创建tensorflow2.0-gpu虚拟环境

为了不影响其他搭建好的编译环境，这里选择重新创建一个虚拟环境，在anaconda中创建虚拟环境的方法见另一篇文章：anaconda中创建虚拟环境

1.2.3 安装CUDA和cudnn

由于模型比较大，需要使用GPU进行训练，使用GPU训练就需要安装CUDA和cudnn，具体的安装方法可以参考我的另一篇文章：CUDA和cudnn的安装

1.2.4 1.4 安装tensorflow GPU版

由于模型比较大，需要在gpu下运行，所以在装tensorflow的时候需要装gpu版本的，由于该源码要求使用tensorflow2.0版本，所以这里我装的是tensorflow-gpu2.3.0
第一步：打开命令提示符
第二步：激活tensorflow2.0-gpu环境（我装的虚拟环境取名为tensorflow2.0-gpu），输入命令：

activate tensorflow2.0-gpu

第三步：安装tensorflow-gpu，输入命令：

pip install  tensorflow-gpu==2.3.0

然后就是等待安装完成，可能需要一定的时间。
第四步：测试tensorflow-gpu是否安装成功，接着上一步，先输入命令打开python编译环境:

python

再输入命令导入tensorflow模块，需要注意的是，虽然装的是tensorflow-gpu，但是我们在导入模块的时候仍然使用的是tensorflow：

import tensorflow as tf

若能成功导入模块，则说明安装成功，否则就需要查看在哪里出了错，重新安装了。
另外还可以输入以下命令来查看tensorflow的版本：

tf.__version__

二、环境测试

在训练自己的数据集之前，先测试一下能不能利用作者训练的好的权重进行而检测。

2.1 下载权重文件

测试之前需要准备YOLOv4权重文件，可以去百度云自行提取，提取密码：o6hd
将下载的权重文件yolov4.weights放到tensorflow-yolov4-tflite-master\data文件夹下
在这里插入图片描述

2.2 权重文件转换

在这一步需要将前面下载的权重文件yolov4.weights转换成在tensorflow环境下使用的文件类型，打开命令提示符，进入主目录（tensorflow-yolov4-tflite-master\）：
在这里插入图片描述
输入下面的命令：

python save_model.py --weights ./data/yolov4.weights --output ./checkpoints/yolov4-416 --input_size 416 --model yolov4

运行结束后，将会在tensorflow-yolov4-tflite-master\checkpoints文件夹下产生一个yolov4-416文件夹
在这里插入图片描述
打开yolov4-416文件夹，里面会有一个saved_model.pb文件，这个就是我们进行测试需要使用的权重文件了。
在这里插入图片描述

2.3 测试

（1）图片测试。打开命令提示符，进入主目录（tensorflow-yolov4-tflite-master\），输入以下命令：

python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --image ./data/kite.jpg --output ./data/kite_d.jpg

就能得到如下测试结果
在这里插入图片描述
并且在tensorflow-yolov4-tflite-master\data文件夹下将产生一个kite_d.jpg文件，也就是我们保存的带有检测结果的图片。
在这里插入图片描述
（2）视频测试。打开命令提示符，进入主目录（tensorflow-yolov4-tflite-master\），输入以下命令：

python detectvideo.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --video ./data/road.mp4 --output ./data/road_d.mp4

测试结果如下：
在这里插入图片描述
并且在tensorflow-yolov4-tflite-master\data文件夹下将产生一个road_d.mp4文件，也就是我们保存的带有检测结果的视频。
在这里插入图片描述

三、使用自己的数据训练

改源码的作者在代码下方说了，该代码的“训练性能还没有完全复制出来，所以我建议使用Alex的Darknet来训练你自己的数据，然后把.weights转换成tensorflow或tflite。”
在这里插入图片描述
不知道现在这版代码是否已经将训练性能完全复制出来了，所以保险起见，还是建议使用YOLOv4原文作者Alex给出的Darknet源码进行训练。
使用Darknet源码进行训练的操作步骤可以参考我上一篇文章：Windows10+YOLOv4+Darknet训练自己的数据（二）
这里就不再复述使用Darknet进行训练的过程了。

下面还是讲述一下使用本文这个tensorflow版本的代码进行训练的步骤。

3.1 数据集制作

关于数据集的制作，大家可以参考我的另一篇文章：VOC数据集制作
这里就不再特别讲述了。

3.2 在data/dataset文件夹下生成数据的描述文件

根据作者在github上的声明，该程序在训练时还需要两个对图片信息的描述文件voc_train.txt和voc_test.txt，格式如下
在这里插入图片描述
作者给出了生成这两个文件的代码scripts\voc_annotation.py。需要做部分的修改：
在这里插入图片描述
（1）将第7行的classes修改为自己的类别名称：

classes = ['类别1','类别2','类别3','类别4']

（2）将相关路径修改为自己的路径
我的数据集完整路径如下，大家可以参考，然后根据自己的数据集修改下面的路径
在这里插入图片描述

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 将相关路径中default的路径修改为自己的路径，最好是写绝对路径，不然后面容易出错
    parser.add_argument("--data_path", default="E:/tensorflow-yolov4-tflite-master/VOCdevkit/")
    # 将default的路径更改为自己的voc_train.txt存放的位置
    parser.add_argument("--train_annotation", default="../data/dataset/voc_train.txt")
     # 将default的路径更改为自己的voc_test.txt存放的位置
    parser.add_argument("--test_annotation",  default="../data/dataset/voc_test.txt")
    flags = parser.parse_args()

    if os.path.exists(flags.train_annotation):os.remove(flags.train_annotation)
    if os.path.exists(flags.test_annotation):os.remove(flags.test_annotation)

    # 更改训练集和测试集的相对路径，由于本实验的训练数据只有一个路径，所以注释掉了num2
    num1 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTrainVal'), 'trainval', flags.train_annotation, False)
    #num2 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTrainVal'), 'trainval', flags.train_annotation, False)
    num3 = convert_voc_annotation(os.path.join(flags.data_path, 'VOCTest'),  'test', flags.test_annotation, False)
    print('=> The number of image for train is: %d\tThe number of image for test is:%d' %(num1, num3))

做了以上更改后直接运行该文件即可。这时候data\dataset文件加下将多出两个文件：
在这里插入图片描述

3.3 修改其他文件

3.3.1 修改voc_names

在tensorflow-yolov4-tflite-master\data\classes文件夹下有个voc.name文件
在这里插入图片描述
打开它，将里面的类别名称就改为自己的类别名称，一个类别名占一行。
在这里插入图片描述

3.3.2 修改配置文件

打开tensorflow-yolov4-tflite-master\core文件夹下的config.py文件
在这里插入图片描述
主要就是修改里面的几个文件的路径，如果是按照我上面步骤得到的那几个文件，那么将不需要修改，如果你是将上面几个文件放在了其他路径下，则需要做相应的修改。
（1）将第14行的“__C.YOLO.CLASSES”修改为自己的voc.names所在位置

__C.YOLO.CLASSES              = "./data/classes/voc.names"

（2）将第29行“__C.TRAIN.ANNOT_PATH”修改为自己的voc_train.txt所在位置

__C.TRAIN.ANNOT_PATH          = "./data/dataset/voc_train.txt"

（3）将第45行“__C.TEST.ANNOT_PATH”修改为自己的voc_test.txt所在位置

__C.TEST.ANNOT_PATH           = "./data/dataset/voc_test.txt"

（4）第30行和第46行的“__C.TRAIN.BATCH_SIZE”和“__C.TEST.BATCH_SIZE”根据自己电脑的性能修改，性能好就调大点，性能差就调小点，没有明确的要求。
config.py文件中的其他参数也是一样的，可以根据需要自行调节，都没有明确的要求。

3.4 训练

3.4.1 从头开始训练

（1）将config.py中第37行的“__C.TRAIN.FISRT_STAGE_EPOCHS”设置为0

__C.TRAIN.FISRT_STAGE_EPOCHS    = 0

（2）直接运行主文件夹下的train.py
在这里插入图片描述
或者在主目录下打开命令提示符，输入命令：

python train.py

3.4.2 从yolov4.weight进行迁移学习训练

或者在主目录下打开命令提示符，输入命令：

python train.py --weights ./data/yolov4.weights

四、使用自己的数据集测试

注意：不管是使用darknet训练的模型，还是使用tensorflow训练出来的模型进行测试，都需要进执行上一节3.1，3.2，3.3中的操作

由于我还没有使用上述的训练方法进行训练，所以这一部分的测试先放一放，以后再补充，现在主要说一下利用darknet训练出来的模型进行测试

4.1 利用darknet训练出来的模型进行测试

将darknet训练出来的模型yolov4-obj_last.weights复制到tensorflow-yolov4-tflite-master\data文件夹下
在这里插入图片描述

4.1.1 测试mAP

在主目录下有一个evaluate.py文件，我们主要需要用到这个文件。
在这里插入图片描述
（1）权重文件转换。打开命令提示符，进入主目录（tensorflow-yolov4-tflite-master\）
在这里插入图片描述
输入下面的命令：

python save_model.py --weights ./data/yolov4-obj_last.weights --output ./checkpoints/yolov4-416 --input_size 416 --model yolov4

运行结束后，将会在tensorflow-yolov4-tflite-master\checkpoints文件夹下产生一个yolov4-416文件夹
（2）修改evaluate.py文件。将文件第20行的“annotation_path”的默认值修改为我们自己数据所在路径，也就是3.2小节中生成的voc_test.txt文件所在位置
在这里插入图片描述
（3）在上一步的命令行模式下，输入以下命令：

python evaluate.py --weights ./checkpoints/yolov4-416

也可以直接在pycharm中运行evaluate.py文件。
（4）若类别名中有空格，需要先将空格去除，不然后续计算mAP会出错，在tensorflow-yolov4-tflite-master\mAP\extra文件夹下有一个remove_space.py文件
打开remove_space.py文件，修改第78行的文件名，将with open('../../data/classes/coco.names') as f:改为with open('../../data/classes/voc.names') as f:然后运行这个文件就可以去除类别名中的空格啦。
在刚刚的命令行模式下先输入命令：

cd mAP/extra

切换到tensorflow-yolov4-tflite-master\mAP\extra文件夹中，再输入命令：

python remove_space.py

也可以直接在pycharm中运行remove_space.py文件。
（5）计算mAP。在tensorflow-yolov4-tflite-master\mAP文件夹下有一个main.py文件，运行这个文件就能得到mAP啦
在刚刚的命令行模式下先输入命令;

cd ..

切换到tensorflow-yolov4-tflite-master\mAP文件夹中，再输入命令：

python main.py

也可以直接在pycharm中运行main.py文件。

最后计算出来的mAP特别低，也不知道为啥，但是同一个模型在darknet中计算出来的mAP却非常高。这个问题以后再解决吧

4.1.2 图片测试

打开命令提示符，进入主目录（tensorflow-yolov4-tflite-master\），输入以下命令：

python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --image ./VOCdevkit/VOCTest/JPEGImages/test1.jpg --output ./VOCdevkit/VOCTest/results/test1.jpg

其中：–image后面跟的是测试图片所在位置；–output后面跟的是检测后的图片保存的位置，这两个根据自己的设置。

4.1.3 视频测试

打开命令提示符，进入主目录（tensorflow-yolov4-tflite-master\），输入以下命令：

python detectvideo.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --video ./data/test.mp4 --output ./data/results/test.mp4

其中：–video后面跟的是测试视频所在位置；–output后面跟的是检测后的视频保存的位置，这两个根据自己的设置。

版权声明：本文为qq_40927867原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接和本声明。