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

一、标注数据

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

YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第1张图片

生成的xml文件

二、制作自己的数据集

我这里总共标注训练数据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文件,文件格式如下图
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第2张图片
生成的txt文件和jpg文件放到同一个目录下,这里我放到了./build/darknet/x64/data/obj。

接下来,就是按照readme的操作制作cfg,train.txt,obj.data.obj.names文件。
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第3张图片
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第4张图片
2.1. 创建文件yolo-obj.cfg,修改配置文件
点1 三个yolo层前面的卷积层,修改方法为filters=(classes + 5)x3),例如,我就一类,那么filters为18=(1+5)*3
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第5张图片
2.2 yolo层classes改为1
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第6张图片
2.3 可以根据自己的情况修改batch等一些训练参数
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第7张图片
2.4 在build\darknet\x64\data\内创建obj.names
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第8张图片

2.5 在build\darknet\x64\data\内创建obj.data
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第9张图片
2.6 在build\darknet\x64\data\内创建train.txt
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第10张图片
三、训练
你可以从头开始训练,也可以使用预训练的模型
./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’)”
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第11张图片
六、测试
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第12张图片
YOLOV4模型剪枝后在Xavier开发板上训练自己的数据集_第13张图片

你可能感兴趣的:(yolov4,集成)