目标检测实战(四)：YOLOV4-Tiny 源码训练、测试、验证详细步骤

下载源码和权重文件

源码：https://github.com/AlexeyAB/darknet
权重：https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v4_pre/yolov4-tiny.weights

编译环境

修改makefile(打开darknet目录下makefile文件),
根据具体情况修改

GPU=1        # 使用GPU
CUDNN=1        # 使用GPU
CUDNN_HALF=1        # 混合精度训练，用于加速
OPENCV=1        # 使用opencv
AVX=0
OPENMP=0
LIBSO=1        # 生成libdarknet.so，便于python调用darknet模型
ZED_CAMERA=0
ZED_CAMERA_v2_8=0

#ARCH= -gencode arch=compute_35,code=sm_35 \
#      -gencode arch=compute_50,code=[sm_50,compute_50] \
#      -gencode arch=compute_52,code=[sm_52,compute_52] \
#        -gencode arch=compute_61,code=[sm_61,compute_61]

OS := $(shell uname)

# GeForce RTX 3070, 3080, 3090
# ARCH= -gencode arch=compute_86,code=[sm_86,compute_86]

# Kepler GeForce GTX 770, GTX 760, GT 740
# ARCH= -gencode arch=compute_30,code=sm_30

# Tesla A100 (GA100), DGX-A100, RTX 3080
ARCH= -gencode arch=compute_80,code=[sm_80,compute_80]

# Tesla V100
# ARCH= -gencode arch=compute_70,code=[sm_70,compute_70]

NVCC=/usr/local/cuda-11.1/bin/nvcc

然后直接终端进行编译

sudo make

就会在当前文件夹生成libdarknet.so文件。

简单测试

./darknet detector test cfg/coco.data cfg/yolov4-tiny.cfg yolov4-tiny.weights data/dog.jpg    # 图片测试
./darknet detector demo cfg/coco.data cfg/yolov4-tiny.cfg yolov4-tiny.weights -ext_output test.mp4        # 视频测试
./darknet detector demo cfg/coco.data cfg/yolov4-tiny.cfg yolov4-tiny.weights -c 0        # 摄像头测试

训练VOC数据集

通过下面代码划分，创建一个data_spilt.py文件和对应的路径

import os
import random

trainval_percent = 0.8  # 所有数据中用来训练的比例 trainval/all  trainval=train+val
train_percent = 1     # trainval用来训练的比例 train/trainval
xmlfilepath = '/home/lqs/Downloads/dataset/VOC/VOC2007/Annotations'
txtsavepath = './VOCdevkit/main'
total_xml = os.listdir(xmlfilepath)

num=len(total_xml)
list=range(num)
tv=int(num*trainval_percent)
tr=int(tv*train_percent)
trainval= random.sample(list,tv)
train=random.sample(trainval,tr)

ftrainval = open(txtsavepath+'/trainval.txt', 'w+')    # 训练集数据+验证集数据
ftest = open(txtsavepath+'/test.txt', 'w+')    # 测试集
ftrain = open(txtsavepath+'/train.txt', 'w+')    # 训练集
fval = open(txtsavepath+'/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()
print('Finished!')

生成训练文件

修改sets和classes以及里面的关键路径即可，这个代码可自己创建，得到2007_test.txt、2007_train.txt、2007_val.txt三个文件

#---------------------------------------------#
#   运行前一定要修改classes
#   如果生成的2007_train.txt里面没有目标信息
#   那么就是因为classes没有设定正确
#---------------------------------------------#
import xml.etree.ElementTree as ET
from os import getcwd

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
#-----------------------------------------------------#
#   这里设定的classes顺序要和model_data里的txt一样
#-----------------------------------------------------#
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

def convert_annotation(year, image_id, list_file):
    in_file = open('/home/lqs/Downloads/dataset/VOC/VOC%s/Annotations/%s.xml'%(year, image_id), encoding='utf-8')
    tree=ET.parse(in_file)
    root = tree.getroot()

    for obj in root.iter('object'):
        difficult = 0 
        if obj.find('difficult')!=None:
            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 = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))

wd = getcwd()

for year, image_set in sets:
    image_ids = open('/home/lqs/Downloads/dataset/VOC/VOC%s/ImageSets/Main/%s.txt'%(year, image_set), encoding='utf-8').read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
    for image_id in image_ids:
        list_file.write('/home/lqs/Downloads/dataset/VOC%s/JPEGImages/%s.jpg'%(year, image_id))
        convert_annotation(year, image_id, list_file)
        list_file.write('\n')
    list_file.close()

训练准备

• 修改darknet-master/cfg/voc.data

classes= 20        # 改成自己的类别数
train  = /home/pjreddie/data/2007_train.txt    # 改成voc_label.py生成的2007_train.txt路径
valid  = /home/pjreddie/data/2007_test.txt    # 改成voc_label.py生成的2007_test.txt路径
names = data/voc.names    # 改成有自己类别的names文件
backup = backup/    # 改为backup/即可

• 修改darknet-master/cfg/yolov4-tiny.cfg

[net]
# Testing
#batch=1
#subdivisions=1
# Training
batch=64    # 每64张图片更新一次参数
subdivisions=16    # 64张图片分16次放入显卡中，每次4张。因为一次放入过多会内存不足，根据自己显卡性能更改
width=416    # 将输入的图片resize到width×height后，放入网络中训练
height=416    # 只要是32的倍数即可
channels=3
momentum=0.9
decay=0.0005
angle=0
saturation = 1.5
exposure = 1.5
hue=.1

• 找到如下位置（以yolo为关键字），修改filters和classes，整个文本共有2个filters和2个classes需要修改

[convolutional]
size=1
stride=1
pad=1
filters=255    # 改为3*(classes +5)。不过官网上不建议这种做法，具体可以查官方github 修改地方不止一处！
activation=linear

[yolo]
mask = 3,4,5
anchors = 10,14,  23,27,  37,58,  81,82,  135,169,  344,319
classes=80    # 改为自己的类别数
num=6
jitter=.3
scale_x_y = 1.05
cls_normalizer=1.0
iou_normalizer=0.07

开始训练

在darknet目录下，使用第一条命令生成预训练权重yolov4-tiny.conv.29，第二条命令开始训练

./darknet partial cfg/yolov4-tiny.cfg yolov4-tiny.weights yolov4-tiny.conv.29 29    # 生成yolov4-tiny.conv.29文件，用于迁移学习
./darknet detector train cfg/voc.data cfg/yolov4-tiny.cfg yolov4-tiny.conv.29 -map    # 训练模型
# -map参数可以在训练过程中对测试集计算map并在chart上显示)

Error in load_data_detection() - OpenCV

多GPU训练

./darknet detector train [.data] [.cfg] [.weight] -gpus 0, 1, 2    # 多GPU训练