Google Colab训练yolov3自定义数据集

共分7步：

一、进入Google Colab训练环境
二、下载编译Darknet
三、上传数据集
四、修改配置文件
五、下载预训练权重文件
六、开始训练数据
七、测试训练结果

一、进入Google Colab训练环境

打开 Google Drive 云端硬盘

新建Google Colaboratory文件，进入python运行环境。 

二、下载编译Darknet

yolov3基于darknet深度学习框架。

1. 运行如下代码，连接到Google Drive 云端硬盘

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

 弹出提示框时，按提示点击允许连接即可。

2.下载darknet

import os
path = '/content/drive/MyDrive/darknet'
if not os.path.isdir(path):
  !git clone https://github.com/AlexeyAB/darknet.git

3. 修改makefile，开启GPU、cuda加速

 也可运行命令

%cd /content/drive/MyDrive/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
!sed -i 's/CUDNN_HALF=0/CUDNN_HALF=1/' Makefile
!sed -i 's/LIBSO=0/LIBSO=1/' Makefile

 4.验证CUDA版本 

!/usr/local/cuda/bin/nvcc --version

 5.编译生成darknet项目。

!make

三、上传数据集

1. 将图片（*.jpg）和对应的标注文件（*.txt）放到同一文件夹下，并将数据放到darknet/data/obj文件夹下。我在网上找的口罩数据集。

 

四、修改配置文件

配置文件共有4个，即obj.data、obj.names、train.txt、yolov3-custom2.cfg。

 2.配置obj.data（位于darknet/data）

obj.data为训练数据的总配置文件。classes为类别数，train、valid、names路径要正确，valid可有可无。backup训练过程中或训练后的权重文件存储路径。

classes= 2
train  = data/train.txt
valid = data/valid.txt
names = data/obj.names
backup = backup/

  3.配置obj.names（位于darknet/data）

obj.names为类别名列表

mask1
mask2

 4.生成train.txt（放到darknet/data） 

 定位到darknet文件夹

%cd /content/drive/
%ls
%cd MyDrive/
%ls
%cd darknet/
%ls

 运行以下代码生成train.txt，注意数据集图片在该路径下（/content/drive/MyDrive/darknet/data/obj）。

import os

image_files = []
os.chdir("/content/drive/MyDrive/darknet/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("..")

 生成的train.txt位于darknet下，需要手动移动到darknet/data下

 5.配置yolov3-custom2.cfg（位于darknet/cfg） 

复制darknet/cfg/yolov3.cfg，重命名为yolov3-custom2.cfg，打开并编辑

 

 修改一

  修改二，搜索yolo，修改classes、filters，共3处。

五、下载预训练权重文件

!wget http://pjreddie.com/media/files/darknet53.conv.74

 

六、开始训练数据

!./darknet detector train data/obj.data cfg/yolov3_custom2.cfg darknet53.conv.74 -dont_show

 等待数小时（我等了3个小时）后训练完成。

查看loss下降过程， 打开darknet/chart.png

权重文件位于darknet/backup 七、测试训练结果

!./darknet detector test data/obj.data cfg/yolov3_custom2.cfg ./backup/yolov3_custom2_last.weights masktest.jpg -thresh 0.3