YOLOX训练测试评估推理加速自己的数据集(一篇博客全详细)

YOLOX学习笔记

环境搭建

conda环境包管理

推荐使用miniconda包，比较轻量级，不需要安装过多的包，conda最主要是管理python环境的工具，我有yolov5、yolov7，yolox三种环境，每个yolo包都拥有专属的环境，可以安装专属的包，环境之间相互独立，互不影响

获取YOLOX源码

 git clone https://github.com/Megvii-BaseDetection/YOLOX

配置YOLOX环境

pip install -r requirements.txt
python3 setup.py develop

注意!!!我们下载pytorch版本的时候尽量不要用官方给的requirements.txt安装，建议注释掉pytorch安装，因为官方给的是安装pytorch最新版本，但是pytorch版本需要严格控制与cuda等版本一致，否则无法使用GPU训练等一系列操作

安装pytorch

需要：pytorch>=1.7
去官网查看对应的CUDA版本号：CUDA版本号，我的CUDA是10.2的

conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 cudatoolkit=10.2

数据集收集

数据集收集

之前收集了不少免费数据集平台，这次选用的是百度飞桨官方的数据集平台

百度飞桨数据集平台

在众多草莓数据集中选择了带有json标签的数据集，因为其他数据集需要自己去瞄框，需要使用一款常用的labelimg数据集标注软件去瞄框得到xml标签文件，我直接选择已有瞄框数据的数据集

数据集格式

首先数据集标签有很多种，csv文件、json文件、xml文件，txt文件

常用的数据集格式标准有两种：VOC和COCO

统一标准有助于神经网络开发工作，每个人都按照标准来制作数据集，那么我们喂进神经网络的数据就可以用写好的代码固定去搜索标准文件夹来获取我们想要的数据

这里采用的是VOC2007数据集，具体格式如下

蓝色的文件夹名称必须按照VOC2007的规定来，方便检索读取文件

ImageSets里面的txt文件名也必须按格式命名，方便检索读取

xml标签文件和jpg原图像的命名可以自定义

• xml标签数据

• txt文件中放置训练集、验证集、测试集的索引

• jpg原图像

数据集处理

这里是通过一些Python基础操作，网上一些参考博客来写的，主要是我自己对json文件的类似字典的读取，对键、对键的对象的遍历进行一点修改和完善

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import os
import numpy as np
import codecs
import json
from glob import glob
import cv2
import shutil
from sklearn.model_selection import train_test_split

#1.标签路径
labelme_path = "../../origin_data/strawberry/"          #原始label标注数据路径
saved_path = "../../datasets/strawberry/VOC2007/"       #保存路径

#2.创建要求文件夹
if not os.path.exists(saved_path + "Annotations"):
    os.makedirs(saved_path + "Annotations")
if not os.path.exists(saved_path + "JPEGImages/"):
    os.makedirs(saved_path + "JPEGImages/")
if not os.path.exists(saved_path + "ImageSets/Main/"):
    os.makedirs(saved_path + "ImageSets/Main/")
    
#3.获取待处理文件
files = glob(labelme_path + "*.json")
files = [i.split("/")[-1].split(".json")[0] for i in files]

#4.读取标注信息并写入 xml
for json_file_ in files:
    json_filename = labelme_path + json_file_ + ".json"
    json_file = json.load(open(json_filename,"r",encoding="utf-8"))
    height, width, channels = cv2.imread(labelme_path + json_file_ +".jpg").shape
    with codecs.open(saved_path + "Annotations/"+json_file_ + ".xml","w","utf-8") as xml:
        xml.write('<annotation>\n')
        xml.write('\t<folder>' + 'UAV_data' + '</folder>\n')
        xml.write('\t<filename>' + json_file_ + ".jpg" + '</filename>\n')
        xml.write('\t<source>\n')
        xml.write('\t\t<database>The UAV autolanding</database>\n')
        xml.write('\t\t<annotation>UAV AutoLanding</annotation>\n')
        xml.write('\t\t<image>flickr</image>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t</source>\n')
        xml.write('\t<owner>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t\t<name>NULL</name>\n')
        xml.write('\t</owner>\n')
        xml.write('\t<size>\n')
        xml.write('\t\t<width>'+ str(width) + '</width>\n')
        xml.write('\t\t<height>'+ str(height) + '</height>\n')
        xml.write('\t\t<depth>' + str(channels) + '</depth>\n')
        xml.write('\t</size>\n')
        xml.write('\t\t<segmented>0</segmented>\n')
        # 获取json文件中的labels对象序列
        # 使用迭代器一个一个取出标签
        for multi in json_file.get("labels"):
            # points = np.array(multi["points"])
            # xmin = min(points[:,0])
            # xmax = max(points[:,0])
            # ymin = min(points[:,1])
            # ymax = max(points[:,1])
            xmin = int((multi['x1']))
            xmax = int((multi['x2']))
            ymin = int((multi['y1']))
            ymax = int((multi['y2']))
            label = multi["name"]
            if xmax <= xmin:
                pass
            elif ymax <= ymin:
                pass
            else:
                xml.write('\t<object>\n')
                xml.write('\t\t<name>'+ str(label)+'</name>\n') 
                xml.write('\t\t<pose>Unspecified</pose>\n')
                xml.write('\t\t<truncated>1</truncated>\n')
                xml.write('\t\t<difficult>0</difficult>\n')
                xml.write('\t\t<bndbox>\n')
                xml.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
                xml.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
                xml.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
                xml.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
                xml.write('\t\t</bndbox>\n')
                xml.write('\t</object>\n')
                print(json_filename,xmin,ymin,xmax,ymax,label)
        xml.write('</annotation>')
        
#5.复制图片到 VOC2007/JPEGImages/下
image_files = glob(labelme_path + "*.jpg")
print("copy image files to VOC007/JPEGImages/")
for image in image_files:
    shutil.copy(image,saved_path +"JPEGImages/")
    
#6.split files for txt
txtsavepath = saved_path + "ImageSets/Main/"
ftrainval = open(txtsavepath+'/trainval.txt', 'w')
ftest = open(txtsavepath+'/test.txt', 'w')
ftrain = open(txtsavepath+'/train.txt', 'w')
fval = open(txtsavepath+'/val.txt', 'w')
total_files = glob(saved_path+"Annotations/*.xml")
total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]
#test_filepath = ""
for file in total_files:
    ftrainval.write(file + "\n")
#test
#for file in os.listdir(test_filepath):
#    ftest.write(file.split(".jpg")[0] + "\n")
#split
train_files,val_files = train_test_split(total_files,test_size=0.15,random_state=42)
#train
for file in train_files:
    ftrain.write(file + "\n")
#val
for file in val_files:
    fval.write(file + "\n")

ftrainval.close()
ftrain.close()
fval.close()
#ftest.close()

训练数据集

修改类别数

修改exps/example/yolox_voc/yolox_voc_s.py

class Exp(MyExp):
  def __init__(self):
      super(Exp, self).__init__()
      self.num_classes = 3 #修改类别数目
      self.depth = 0.33
      self.width = 0.50
      self.warmup_epochs = 1

修改类别名

修改yolox/data/datasets/voc_classes.py为自己的类别

VOC_CLASSES = (
	'mature',
 	'growth',
 	'flower'
)

开始训练

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 4 --fp16 -c ./weights/yolox_s.pth

• 训练中可选参数及其意义如下：

  -expn 或 --experiment-name # 实验的名字 接受str类的参数 默认为None
  -n 或 --name # 模型的名字 接受str类的参数 默认为None
  –dist-backend # 分布式的后端 接受str类的参数 默认为nccl
  –dist-url # 分布式训练的url地址 接受str类的参数 默认为None
  -b 或 --batch-size # 训练批量大小 接受int类的参数 默认为64 实际建议设在16以下 具体取决于显存大小
  -d 或 --devices # 训练的设备（GPU）数量 接受int类的参数 默认为None
  -f 或 --exp_file # 训练的模型声明文件 接受str类的参数 默认为None
  –resume # 是否从上一个checkpoint继续训练，一般中断后继续训练时用，直接输入–resume不需要跟参数
  -c 或 --ckpt # 上次训练的结果，继续训练和fine turning时填写check point路径 默认None 接受str类的参数
  -e 或 --start_epoch # 指定开始的 epoch 和 --resume 等参数搭配使用 接受int类的参数 默认为None
  –num_machines # 分布式训练的机器数 接受int类的参数 默认为1
  –machine_rank # 分布式训练中该机器的权重 接受int类的参数 默认为0
  –fp16 # 在训练时采用混合精度 默认False 只要输入了–fp16就是True 无需参数
  –cache # 是都将图片缓存在RAM，RAM空间够的话推荐开启，能够加速训练 默认False 只要输入了–cache就是True 无需参数
  -o 或 --occupy # 在训练开始时是否占用所有需要的GPU内存 默认False 只要输入了就是True 无需参数
  -l 或 --logger # 记录训练的logger 接受str类的参数 默认为 tensorboard 可选 wandb

如果训练终端，可以重新接着开启训练，使用–resume

python3 tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0 -b 64 -c <last_epoch_ckpt.pth的路径> --resume

测试模型

修改数据集标准

yolox是默认加载coco数据集的，所以我们需要更改为VOC数据集格式

修改yolox/data/datasets/下的init.py文件，添加：

from .voc_classes import VOC_CLASSES

修改tools/demo.py文件，将COCO_CLASSES全部修改为VOC_CLASSES

开始测试

python tools/demo.py image -f exps/example/yolox_voc/yolox_voc_s.py -c weights/best_ckpt.pth --path assets/class01.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]

• 推理中可选参数及其意义如下：

  直接在tools/demo.py后输入文件类型：image video 或者 webcam
  -expn 或 --experiment-name # 实验的名字 接受str类的参数 默认为None
  -n 或 --name # 模型的名字 接受str类的参数 默认为None
  –path # 需要预测的文件夹 默认为./assets/dog.jpg
  –camid # webcam demo camera id 接受int类的参数 默认为0
  –save_result # 是否保存结果 直接输入 不需要跟参数
  -f 或 --exp_file # 训练的模型声明文件 接受str类的参数 默认为None
  -c 或 --ckpt # 上次训练的结果
  –device # 推理的设备 cpu gpu 二选一
  –conf # 输出框的最低置信度 接受float类的参数 默认为0.3
  –nms # nms阈值 接受float类的参数 默认为0.3
  –tsize # 测试图片大小 接受int类的参数 默认为None
  –fp16 # 在训练时采用混合精度 默认False 只要输入了–fp16就是True 无需参数
  –legacy # 配饰旧版本 默认False 只要输入了就是True 无需参数
  –fuse # Fuse conv and bn for testing 默认False 只要输入了就是True 无需参数
  –trt # 在测试时是否使用TensorRT模型 只要输入了就是True 无需参数

推理加速

yolox非常良心，开源的同时让实习生们写了一些推理加速的demo供我们使用

(因为我用VScode进行开发，git history这个插件可以看到提交者哈哈哈哈)

我们目前尝试了两种推理引擎，分别是TensorRT和ONNXRuntime

基本上每个demo都有相应十分详尽的README.md介绍

TensorRT推理加速

我觉得官方的TensorRT的README文档写得很好，上手挺快的，就是配置TensorRT环境比较麻烦，有很详细的英文注释，跟着一步一步操作即可

转换engine模型

python tools/trt.py -f <YOLOX_EXP_FILE> -c <YOLOX_CHECKPOINT>

转换TensorRT模型，后缀名为.engine，这是TensorRT专属的推理加速引擎

Python推理加速

官方对python的优化，把TensorRT的推理加速，已经融入到demo.py的参数表里面，所以我们只要配好环境，只用–trt 就可以很方便调用TensorRT推理加速

python tools/demo.py image -f exps/default/yolox_s.py --trt --save_result

C++推理加速

同样的跟着官方的README走，需要用python导出 .engine引擎文件，就可以了，不再赘述

./yolox <path/to/your/engine_file> -i <path/to/image>

ONNXruntime推理加速

转换onnx模型

python3 tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth

转换为onnx模型，全称**开放神经网络交换(Open Neural Network Exchange, ONNX)**是一种用于表示机器学习模型的开放标准文件格式

python多张图片推理加速

同样的也是阅读官方README文档就可以实现ort推理加速

python3 onnx_inference.py -m <ONNX_MODEL_PATH> -i <IMAGE_PATH> -o <OUTPUT_DIR> -s 0.3 --input_shape 640,640

但是官方的例程不支持多张图片输入同时进行加速推理，每次只能一张进行推理，所以我参考/tools/demo.py的代码，对源码进行一些修改，成功实现多张图片加速推理

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.

import argparse
import os
import time
import cv2
import numpy as np

import onnxruntime

from yolox.data.data_augment import preproc as preprocess
from yolox.data.datasets import VOC_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis

from loguru import logger

IMAGE_EXT = [".jpg", ".jpeg", ".webp", ".bmp", ".png"]


# 新增加的分割路径中的图片名称函数
def get_image_list(path):
    image_names = []
    for maindir, subdir, file_name_list in os.walk(path):
        for filename in file_name_list:
            apath = os.path.join(maindir, filename)
            ext = os.path.splitext(apath)[1]
            if ext in IMAGE_EXT:
                image_names.append(apath)
    return image_names



if __name__ == '__main__':
    args = make_parser().parse_args()

    input_shape = tuple(map(int, args.input_shape.split(',')))
    # origin_img = cv2.imread(args.image_path)
    
    session = onnxruntime.InferenceSession(args.model)


    files = get_image_list(args.image_path)
    files.sort()
    
	# 利用迭代器进行遍历每张图片
    for img_name in files:
        origin_img = cv2.imread(img_name)
        img, ratio = preprocess(origin_img, input_shape)
        
        ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}

        t0 = time.time()
        output = session.run(None, ort_inputs)
        logger.info("Infer time: {:.4f}s".format(time.time() - t0))
        logger.info("img:{}".format(os.path.basename(img_name)))
        predictions = demo_postprocess(output[0], input_shape, p6=args.with_p6)[0]

        boxes = predictions[:, :4]
        scores = predictions[:, 4:5] * predictions[:, 5:]

        boxes_xyxy = np.ones_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
        boxes_xyxy /= ratio
        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
        if dets is not None:
            final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
            origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                            conf=args.score_thr, class_names=VOC_CLASSES)

        mkdir(args.output_dir)
        output_path = os.path.join(args.output_dir, os.path.basename(img_name))
        cv2.imwrite(output_path, origin_img)

所以更改之后我们输入-i 就可以输入路径和图片，真正实现多张图片加速推理

模型评估

在目标检测领域中，我们经常对模型的性能进行评估，有很多性能参数，目前我对模型的推理时间和mAP进行测试，所以有这方面的工作，更多的性能参数有待深入

同样的yolox默认是coco数据集，但是VOC数据集的适配对于模型评估eval特别不好，需要对整个评估的代码进行研读，而且网上的介绍特别少，基本上没有

我是先看eval代码的参数列表，根据demo.py的形式，再根据报错一步一步完善过来的

修改基础模型

修改/yolox/exp/yolo_base.py文件

# 131行左右，获取数据集路径
    from yolox.data import (
        VOCDetection,
        TrainTransform,
        YoloBatchSampler,
        DataLoader,
        InfiniteSampler,
        MosaicDetection,
        worker_init_reset_seed,
    )
    from yolox.utils import wait_for_the_master

    with wait_for_the_master():
        dataset = VOCDetection(
            data_dir=self.data_dir,
            json_file=self.train_ann,
            img_size=self.input_size,
            preproc=TrainTransform(
                max_labels=50,
                flip_prob=self.flip_prob,
                hsv_prob=self.hsv_prob),
            cache=cache_img,
        )

# 273行左右，评估集加载  
    	from yolox.data import VOCDetection, ValTransform

        # valdataset = COCODataset(
        #     data_dir=self.data_dir,
        #     json_file=self.val_ann if not testdev else self.test_ann,
        #     name="val2017" if not testdev else "test2017",
        #     img_size=self.test_size,
        #     preproc=ValTransform(legacy=legacy),
        # )

        valdataset = VOCDetection(
            data_dir=self.data_dir,
            img_size=self.test_size,
            preproc=ValTransform(legacy=legacy),
        )

#  获取yolox评估器
    def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
        from yolox.evaluators import VOCEvaluator

        val_loader = self.get_eval_loader(batch_size, is_distributed, testdev, legacy)
        evaluator = VOCEvaluator(
            dataloader=val_loader,
            img_size=self.test_size,
            confthre=self.test_conf,
            nmsthre=self.nmsthre,
            num_classes=self.num_classes,
        )

注意！！！VOC的数据加载是VOCDetection类，COCO数据加载是COCODatasets类，而且两个类的需要形参都不一样，需要找到这些数据加载类的定义，对应去修改加载相应的参数

• yolox/data/datasets/voc.py

  class VOCDetection(Dataset):
      def __init__(
          self,
          data_dir,
          image_sets=[("2007", "trainval")],
          img_size=(416, 416),
          preproc=None,
          target_transform=AnnotationTransform(),
          dataset_name="VOC2007",
          cache=False,
      ):
  

• yolox/data/datasets/coco.py

  class COCODataset(Dataset):
      def __init__(
          self,
          data_dir=None,
          json_file="instances_train2017.json",
          name="train2017",
          img_size=(416, 416),
          preproc=None,
          cache=False,
      ):
  

注意！！！

每次修改完模型参数之后，需要对yolox模型进行重新加载，相当于重新导入yolox包

在顶层YOLOX文件夹终端下输入

python setup.py install        

开始评估模型

python3 tools/eval.py -f exps/default/yolox_s.py -c YOLOX_outputs/yolox_voc_s/epoch_300_ckpt.pth -d 0 -b 8 --conf 0.001 --fp16 --fuse  

总结

模型结果

推理时间

最近把yolox的全流程跑通一遍，虽然没有全部推理加速引擎都跑试过一遍，感觉有需求才去深入学习，对比各种引擎的优缺点来对比，参考博客的推理部署速度比较，并且因为我电脑有gpu所以先学TensorRT推理加速

(OpenVINO因为我电脑是AMD，所以跑不了，有空找Intel电脑跑)

电脑配置：

CPU AMD R7 显卡 1650

目前推理加速情况：

• cpu推理：0.2s一张
• gpu推理：0.02s一张
• ort推理加速：0.09s一张
• TensorRT推理加速：0.007s一张

模型评估

使用yolox官方给出的eval.py，并且根据VOC2007数据集格式进行修改，迭代了300个epoch时

map_5095: 77.3%

map_50: 90.8%

问题

1. 网络训练出来，background大部分是绿色的，红色对比度过高，导致权重置信度过高，一半红色一半白色的草莓也当成了mature

   通过查看label标签json文件发现，标签有问题，x、y相似的瞄框有两个标签，分别是mature和growth，所以训练出来还是mature的置信度高，并没有被召回
   

2. 网络结构的优化部分，不知道从何下手比较好，从上年论文入手，我们的部署平台是NUC，所以没有gpu只能用cpu，推理时间需要通过后剪枝，数据预处理等解决

网络结构优化完善

对yolov5有很多网络完善方式，网上也有很多教程，对于yolox的可能不多，但是我觉得可以移植过来，yolov7也有部分是参考yolov5的优化方式

这里简单抛砖引玉一下，一种数据增强方式

Mosaic数据增强方法

mosaic数据增强则利用了四张图片，对四张图片进行拼接，每一张图片都有其对应的框框，将四张图片拼接之后就获得一张新的图片，同时也获得这张图片对应的框框，然后我们将这样一张新的图片传入到神经网络当中去学习，相当于一下子传入四张图片进行学习了。论文中说这极大丰富了检测物体的背景！且在标准化BN计算的时候一下子会计算四张图片的数据！

不足

目前正在对yolox的改进方向进行学习，目前只是对yolox的框架有个大概的了解，对于yolox的文件结构感觉get不到文件的逻辑性，相对于yolov5来说，我是目前觉得有点乱，可能需要进一步深入学习，一整个流程跑通一次，方便之后对yolo的网络进行改进，获取更好的准确率和推理时间

参考CSDN博客：

配置Pytorch指定版本环境

json2VOC.py代码

train.py参数

yolox train自定义数据集修改代码

yolox数据增强

YOLOX在OpenVINO、ONNXRUNTIME、TensorRT上面推理部署与速度比较