YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集

一、标注数据

使用labelImg进行标注，标注方法自行研究，标注完成后生成的是xml后缀的文件，该格式是VOC的标准格式。

二、制作自己的数据集

我这里总共标注训练数据4351张，类别为船（共一类），数据集制作参考yoloV4里面的readme。

import os
import xml.etree.ElementTree as ET

if __name__ == '__main__':

    list_file = os.listdir('F:\\155dataSet\\boatAll\\annotations\\')
    for i in list_file:
        annotation_file = f"F:\\155dataSet\\boatAll\\annotations\\{i}"
        txtfile = f"F:\\155dataSet\\boatAll\\labels\\{i}"
        txt_file = f"{txtfile[0:-4]}.txt"
        objects = ET.parse(annotation_file).findall("object")
        sizes = ET.parse(annotation_file).findall("size")
        for size in sizes:
            width = float(size.find('width').text)
            height = float(size.find('height').text)
        for object in objects:
            bbox = object.find('bndbox')
            xmin = float(bbox.find('xmin').text) - 1
            ymin = float(bbox.find('ymin').text) - 1
            xmax = float(bbox.find('xmax').text) - 1
            ymax = float(bbox.find('ymax').text) - 1
            classname = object.find('name').text

            xcenter = (xmax + xmin) / (2 * width)
            ycenter = (ymax + ymin) / (2 * height)
            w = (xmax - xmin) / (width)
            h = (ymax - ymin) / (height)
            if classname == 'boat':
                object_class = 0
            else:
                object_class = 1
            data = [str(object_class),' ',str(xcenter),' ',str(ycenter),' ',str(w),' ',str(h)]

            with open(txt_file,'a+') as f:
                f.writelines(data)
                f.write('\n')

该部分代码适用于xml转化为txt文件，文件格式如下图

生成的txt文件和jpg文件放到同一个目录下，这里我放到了./build/darknet/x64/data/obj。

接下来，就是按照readme的操作制作cfg，train.txt,obj.data.obj.names文件。


2.1. 创建文件yolo-obj.cfg，修改配置文件
点1 三个yolo层前面的卷积层，修改方法为filters=(classes + 5)x3)，例如，我就一类，那么filters为18=（1+5）*3

2.2 yolo层classes改为1

2.3 可以根据自己的情况修改batch等一些训练参数

2.4 在build\darknet\x64\data\内创建obj.names

2.5 在build\darknet\x64\data\内创建obj.data

2.6 在build\darknet\x64\data\内创建train.txt

三、训练
你可以从头开始训练，也可以使用预训练的模型
./darknet detector train data/obj.data yolo-obj.cfg yolov4.conv.137
四、剪枝
剪枝操作我研究不深，完全按照github中的yolov3-channel-and-layer-pruning，使用readme中的稀疏和剪枝操作，对训练好的模型进行剪枝，得到best.pt文件，具体的剪枝方法可以参考readme，里面非常详细。
python slim_prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --global_percent 0.8 --layer_keep 0.01
五、微调训练
我先将pt格式的文件转化为了weights类型的文件，该文件比较容易在darknet中进行微调训练，即我将该weights文件和对应的cfg文件放到darknet中，重新微调训练，训练时可以调整一下参数和训练次数等。
python3 -c “from models import *; convert(‘cfg/prune_0.8_keep_0.01_yolov4.cfg’, ‘weights/best.pt’)”

六、测试