yolov4用Google Colab训练自己的数据集

简介

前段时间YOLO系列(v1-v3)作者 Joe Redmon在推特上宣布不再继续CV方向的研究，但最近由 Alexey Bochkovskiy大神接管并更新了yolov4的项目，项目地址https://github.com/AlexeyAB/darknet。此项目之前为yolov3项目支持windows的版本，由于官方darknet并不能很好得支持windows环境，便采用了支持windows的此项目，并且此项目一直在更新，用起来比大多数yolov3项目都顺手。

对于yolov4部分的训练流程和测试，同yolov3的过程。
yolov3及标注生成方法可参考：https://blog.csdn.net/longlong068/article/details/105384712
https://blog.csdn.net/longlong068/article/details/105088304

运用google Colab训练自己的Yolov4数据集

google colab的使用

1. 首先需要创建谷歌云盘账号：https://drive.google.com

2. 打开谷歌云盘

   

   右键选择创建一个Colab的.ipynb文件，如果没有上图的Googel Colab选项，点击关联更多应用添加。

3. 将.ipynb文件装载谷歌云盘，方便Googel云盘和Colab分配的云GPU进行文件相互传输。

   

   点击装载Google云端硬盘，或者直接在左边的命令行输入指令进行装载，指令如下：

   from google.colab import drive
   drive.mount('/content/drive')
   

   装载成功后可见：

   

   其中的My Drive中的文件即为Google云盘中的文件。

4. 在生成的ipynb文件中进行设置GPU可用、项目的git、训练等后续操作。

项目使用流程

此部分的Yolov4项目基于https://github.com/AlexeyAB/darknet
。

1.设置GPU可用

2.clone项目构建Darknet

git clone项目

输入并运行以下命令来git项目：

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

成功后在左边的文件列表可以找到项目文件darknet。

修改项目的makefile文件，使编译生成的程序可以使用GPU和OpenCV

运行以下命令，将会自动修makefile

# 修改makefile 将OpenCV和GPU设置为可用
%cd darknet
!sed -i 's/OPENCV=0/OPENCV=1/' Makefile
!sed -i 's/GPU=0/GPU=1/' Makefile
!sed -i 's/CUDNN=0/CUDNN=1/' Makefile

验证CUDA版本

#验证CUDA版本
!/usr/local/cuda/bin/nvcc --version

输出结果如下：

编译项目生成darknet运行程序

!make

编译生成可执行文件darknet

3.设置谷歌云盘路径简写——方便上传文件到云端

装载谷歌云盘——便于将谷歌云盘中的文件传到云端服务器中

因为服务器与云盘的装载经常掉，前面装载过也最好再运行以下代码

from google.colab import drive
drive.mount('/content/gdrive')

设置谷歌云盘路径的简写

#将谷歌云盘路径简写为mydrive
!ln -s /content/gdrive/My\ Drive/ /mydrive
!ls /mydrive

运行以上代码，以后输入mydrive即代表谷歌云盘的路径

在谷歌云盘中创建存放训练相关文件的文件夹

回到谷歌云盘中创建一个文件夹，来存放我们训练所要用到的文件，我创建了一个名为yolov4的文件夹，来存放obj.name、obj.zip(训练图片和.txt的压缩文件)等训练所需文件。

!ls /mydrive/yolov4

4.运行demo显示bbox

此步只为了检验环境和编译成功与否，可以跳过此步直接训练自己的数据集。

下载训练好的coco数据集权重

yolov4训练好的coco数据集下载地址：https://drive.google.com/open?id=1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT

也可以通过我的百度云盘下载：链接：https://pan.baidu.com/s/1L5OMQCWcKc76stKj1uaiGQ

提取码：cqhx

要将下载好的权重文件上传到云端的darknet文件下，我是将权重文件上传到Google云盘中的yolov4(自己创建的)文件中，再将其从谷歌云盘copy到云端，copy指令执行以下：

#将谷歌云盘中的yolov4上传到darknet目录下
!cp /mydrive/yolov4/yolov4.weights ./

定义imshow函数，方便显示图片

#定义imshow 调用opencv显示图片
def imShow(path):
  import cv2
  import matplotlib.pyplot as plt
  %matplotlib inline

  image = cv2.imread(path)
  height, width = image.shape[:2]
  resized_image = cv2.resize(image,(3*width, 3*height), interpolation = cv2.INTER_CUBIC)

  fig = plt.gcf()
  fig.set_size_inches(18, 10)
  plt.axis("off")
  plt.imshow(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB))
  plt.show()

运行demo

!./darknet detector test cfg/coco.data cfg/yolov4.cfg yolov4.weights  data/dog.jpg

imShow('predictions.jpg')

结果图：

5.训练自己的数据集

上传训练所需文件

训练所需文件：obj.zip（图片及对应标注的.txt文件压缩包）、yolov4_custom.cfg(自己根据数据集中类修改)、obj.data(相关路径及种类个数设定)、obj.names(自己数据集的类名)、generate_trian.py(运行生成train.txt)

将相关文件上传到谷歌云盘

对自己数据集的图片进行标注和生成.txt文件放到同一文件夹后，将其压缩然后上传到谷歌云盘中，直接上传将十分缓慢。

标注工具：https://github.com/tzutalin/labelImg

上传obj.zip到谷歌云盘：通过标注工具生成项目所需的与图片同名的.txt文件，将原图片与同名.txt文件放在同一文件obj下，将文件夹压缩为obj.zip并上传到谷歌云盘的yolov4文件夹下。

上传obj.data和obj.names到谷歌云盘:obj.data文件在项目中文件的基础上，根据自己数据集的类别个数进行修改而得到。

obj.names中每行写入自己数据集的类名。

上传yolov4_custom.cfg到谷歌云盘：yolov4_custom.cfg在darknet/cfg/yolov4-custom.cfg的基础上，根据自己数据集的类别个数进行修改得到。

修改1：

• batch =64 subdivisions=16不变
  当训练运行显示“out of memory”时，可以尝试将subdivisions升至32或64。
• max_batches= ? 为自己数据集种类个数*2000 如果有六类及为12000
• steps=？，？ 分别为max_batches的80%和90%

修改2：

要修改三个以上部分，即三个尺度下都要修改

• fiters=(5+classes)x3 如数据集种类个数classes为6 fiters=33
• classes =6（分类个数）

上传generate_trian.py文件到谷歌云盘：此yolov4项目训练自己数据集还需要trian.txt，train.txt文件可通过运行generate_trian.py生成，所以在此之前需要创建generate_trian.py，并上传到谷歌云盘中，以便后续拷贝到云服务中并运行。generate_trian.py文件如下：

import os

image_files = []
os.chdir(os.path.join("data", "obj"))
for filename in os.listdir(os.getcwd()):
    if filename.endswith(".jpg"):
        image_files.append("data/obj/" + 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("..")

将云盘中的相关文件拷贝到云端

#确保当前路径在darknet下
!ls
#将数据集图片及标注的txt文件的压缩包上传到服务器
!cp /mydrive/yolov4/obj.zip ../
#解压到data文件下
!unzip ../obj.zip -d data/

#从谷歌云盘上传根据自己数据集修改的.cfg文件
!cp /mydrive/yolov4/yolov4_custom.cfg ./cfg

#从谷歌云盘上传自己数据集的obj.data obj.names
!cp /mydrive/yolov4/obj.names ./data
!cp /mydrive/yolov4/obj.data  ./data

#上传generate_train.py 以在服务器的data下生成train.txt
!cp /mydrive/yolov4/generate_train.py ./

#运行在data文件下生成train.txt
!python generate_train.py

下载预训练权重并上传到darknet文件下

yolov4的预训练权重下载地址：https://drive.google.com/open?id=1JKF-bdIklxOOVy-2Cr5qdvjgGpmGfcbp

也可在我的百度云盘中下载：链接：https://pan.baidu.com/s/1fTxMkvNLEafXmeWXQ47EmA

提取码：c21x

# 介于预训练权重网址为云盘 采用上传方式
!cp /mydrive/yolov4/yolov4.conv.137 ./

开始训练

!./darknet detector train data/obj.data cfg/yolov4_custom.cfg yolov4.conv.137 -dont_show

训练过程中产生的权重文件将保存到darknet/backup下，loss图保存为darknet/chart.png:

如果训练时因不明原因中断，此后想要在上次训练权重的基础上继续训练可采用以下代码

!./darknet detector train data/obj.data cfg/yolov4_custom.cfg backup/yolov4_custom_last.weights  -dont_show

根据map选择权重文件

需要将obj.data中的验证集即valid改为自己验证集生成的.txt文件，同train.txt生成方法相同，如果没有验证集就采用训练集。

将验证集图片和标注信息放到data/test文件下，我是从谷歌云盘拷贝到服务端，如下

#拷贝test.zip到云端
!ls /mydrive
!cp /mydrive/yolov4/image/test.zip  ../
#解压到data文件下
!mkdir  data/test
!unzip ../test.zip -d data/test

以上命令需要根据自己验证集路径，最后只要在data/test生成如下形式即可：

#上传generate_test.py到darknet文件下  以在服务器的data下生成test.txt
!cp /mydrive/olov4/generate_test.py ./
!python generate_train.py

generate_test.py：

import os

image_files = []
os.chdir(os.path.join("data", "test"))
for filename in os.listdir(os.getcwd()):
    if filename.endswith(".jpg"):
        image_files.append("data/test/" + filename)
os.chdir("..")
with open("test.txt", "w") as outfile:
    for image in image_files:
        outfile.write(image)
        outfile.write("\n")
    outfile.close()
os.chdir("..")

如果想要在训练时可以直观看map，可以命令行加-map，采用以下：

! ./darknet detector train data/obj.data cfg/yolov4_custom.cfg yolov4.conv.137 -dont_show -map

其中推荐将权重文件保存在自己的谷歌云盘yolo4/backup/文件下，只需要修改obj.data文件中的backup后面的文件位置，如下：

实际在训练中加-map可能导致出错，可以选择训练不加-map，训练结束后再采用如下指令计算权重文件map：

#不加-iou_thresh 默认计算map@50，即iou=0.5
!./darknet detector map data/obj.data cfg/yolov4_custom.cfg /mydrive/yolov4/backup/yolov4_custom_final.weights
#map@75
!./darknet detector map data/obj.data cfg/yolov4_custom.cfg /mydrive/yolov4/backup/yolov4_custom_final.weights -iou_thresh 0.75

6.测试

上传测试图片到data/obj

#图片测试
!./darknet detector test data/obj.data cfg/yolov4_custom.cfg /mydrive/yolov4/backup/yolov4_custom_final.weights ./data/obj/xxxx.jpg -thresh 0.5
imShow('predictions.jpg')

视频、摄像头等测试见源码GitHub。