win10安装yolox，训练自定义模型，使用tensorrt部署全流程

系统环境：win10、cuda10.2、cudnn8.2

一、采集数据

有2段视频，先使用ffmpeg对视频进行抽帧，由于视频比较长，所以每隔5秒抽取1张图片。

ffmpeg -i light000.mp4 -q:v 2 -vf fps=0.2 ./1%05d.jpg

-q:v 2 表示输出的图片质量，一般是1到5之间（1 为质量最高）

fps为1时按一秒一帧切分，按5秒一帧来切，1/5 = 0.2。

由于一共2个视频，所以我直接采用ffmpeg进行生成，如果你的视频文件比较多，也可以写个python脚本实现抽帧。

二、数据集标注

使用labelme对生成的图片进行标注。

pip install labelme

使用pip安装labelme，然后执行

labelme

打开标注工具，点击“OpenDir”打开图片所在文件夹。

点击File，选择Save Automatically，标记后自动保存，方便标记。 

 

点击Edit，Cteate Rectangle， 

 用矩形框框选目标，标记点输入“dog”，

 将图片全部标注完成后，可以在图片目录下看到每张图片对应一个json文件。

三、labelme数据集转换

yolox需要coco格式的数据集，使用labelme标注的结果不满足要求，所以需要将labelme数据集转换为coco格式。

创建images文件夹，将labelme标注的图片和json文件复制到文件夹下。

创建文件labels.txt，文件内容如下：

__ignore__
dog

由于我只对狗进行分类测试，所以填写dog即可。

创建labelme2coco.py，该脚本将用于数据集转换：

# 命令行执行：  python labelme2coco.py --input_dir images --output_dir coco --labels labels.txt
# 输出文件夹必须为空文件夹

import argparse
import collections
import datetime
import glob
import json
import os
import os.path as osp
import sys
import uuid
import imgviz
import numpy as np
import labelme
from sklearn.model_selection import train_test_split

try:
    import pycocotools.mask
except ImportError:
    print("Please install pycocotools:\n\n    pip install pycocotools\n")
    sys.exit(1)


def to_coco(args,label_files,train):

    # 创建 总标签data 
    now = datetime.datetime.now()
    data = dict(
        info=dict(
            description=None,
            url=None,
            version=None,
            year=now.year,
            contributor=None,
            date_created=now.strftime("%Y-%m-%d %H:%M:%S.%f"),
        ),
        licenses=[dict(url=None, id=0, name=None,)],
        images=[
            # license, url, file_name, height, width, date_captured, id
        ],
        type="instances",
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
        categories=[
            # supercategory, id, name
        ],
    )

    # 创建一个 {类名 : id} 的字典，并保存到 总标签data 字典中。
    class_name_to_id = {}
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()   # strip() 方法用于移除字符串头尾指定的字符（默认为空格或换行符）或字符序列。
        if class_id == -1:
            assert class_name == "__ignore__"   # background:0, class1:1, ,,
            continue
        class_name_to_id[class_name] = class_id
        data["categories"].append(
            dict(supercategory=None, id=class_id, name=class_name,)
        )


    if train:
        out_ann_file = osp.join(args.output_dir, "annotations","instances_train2017.json")
    else:
        out_ann_file = osp.join(args.output_dir, "annotations","instances_val2017.json")


    for image_id, filename in enumerate(label_files):

        label_file = labelme.LabelFile(filename=filename)
        base = osp.splitext(osp.basename(filename))[0]      # 文件名不带后缀
        if train:
            out_img_file = osp.join(args.output_dir, "train2017", base + ".jpg")
        else:
            out_img_file = osp.join(args.output_dir, "val2017", base + ".jpg")
        
        print("| ",out_img_file)


        # ************************** 对图片的处理开始 *******************************************
        # 将标签文件对应的图片进行保存到对应的 文件夹。train保存到 train2017/ test保存到 val2017/
        img = labelme.utils.img_data_to_arr(label_file.imageData)   # .json文件中包含图像，用函数提出来
        imgviz.io.imsave(out_img_file, img)     # 将图像保存到输出路径

        # ************************** 对图片的处理结束 *******************************************

        # ************************** 对标签的处理开始 *******************************************
        data["images"].append(
            dict(
                license=0,
                url=None,
                file_name=base+".jpg",              # 只存图片的文件名
                # file_name=osp.relpath(out_img_file, osp.dirname(out_ann_file)),  # 存标签文件所在目录下找图片的相对路径

                ##   out_img_file : "/coco/train2017/1.jpg"
                ##   out_ann_file : "/coco/annotations/annotations_train2017.json"
                ##   osp.dirname(out_ann_file) : "/coco/annotations"
                ##   file_name : ..\train2017\1.jpg   out_ann_file文件所在目录下 找 out_img_file 的相对路径
                height=img.shape[0],
                width=img.shape[1],
                date_captured=None,
                id=image_id,
            )
        )

        masks = {}  # for area
        segmentations = collections.defaultdict(list)  # for segmentation
        for shape in label_file.shapes:
            points = shape["points"]
            label = shape["label"]
            group_id = shape.get("group_id")
            shape_type = shape.get("shape_type", "polygon")
            mask = labelme.utils.shape_to_mask(
                img.shape[:2], points, shape_type
            )

            if group_id is None:
                group_id = uuid.uuid1()

            instance = (label, group_id)

            if instance in masks:
                masks[instance] = masks[instance] | mask
            else:
                masks[instance] = mask

            if shape_type == "rectangle":
                (x1, y1), (x2, y2) = points
                x1, x2 = sorted([x1, x2])
                y1, y2 = sorted([y1, y2])
                points = [x1, y1, x2, y1, x2, y2, x1, y2]
            else:
                points = np.asarray(points).flatten().tolist()

            segmentations[instance].append(points)
        segmentations = dict(segmentations)

        for instance, mask in masks.items():
            cls_name, group_id = instance
            if cls_name not in class_name_to_id:
                continue
            cls_id = class_name_to_id[cls_name]

            mask = np.asfortranarray(mask.astype(np.uint8))
            mask = pycocotools.mask.encode(mask)
            area = float(pycocotools.mask.area(mask))
            bbox = pycocotools.mask.toBbox(mask).flatten().tolist()

            data["annotations"].append(
                dict(
                    id=len(data["annotations"]),
                    image_id=image_id,
                    category_id=cls_id,
                    segmentation=segmentations[instance],
                    area=area,
                    bbox=bbox,
                    iscrowd=0,
                )
            )
        # ************************** 对标签的处理结束 *******************************************

        # ************************** 可视化的处理开始 *******************************************
        if not args.noviz:
            labels, captions, masks = zip(
                *[
                    (class_name_to_id[cnm], cnm, msk)
                    for (cnm, gid), msk in masks.items()
                    if cnm in class_name_to_id
                ]
            )
            viz = imgviz.instances2rgb(
                image=img,
                labels=labels,
                masks=masks,
                captions=captions,
                font_size=15,
                line_width=2,
            )
            out_viz_file = osp.join(
                args.output_dir, "visualization", base + ".jpg"
            )
            imgviz.io.imsave(out_viz_file, viz)
        # ************************** 可视化的处理结束 *******************************************

    with open(out_ann_file, "w") as f:  # 将每个标签文件汇总成data后，保存总标签data文件
        json.dump(data, f)


# 主程序执行
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--input_dir", help="input annotated directory")
    parser.add_argument("--output_dir", help="output dataset directory")
    parser.add_argument("--labels", help="labels file", required=True)
    parser.add_argument("--noviz", help="no visualization", action="store_true")
    args = parser.parse_args()

    if osp.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    print("| Creating dataset dir:", args.output_dir)
    if not args.noviz:
        os.makedirs(osp.join(args.output_dir, "visualization"))

    # 创建保存的文件夹
    if not os.path.exists(osp.join(args.output_dir, "annotations")):
        os.makedirs(osp.join(args.output_dir, "annotations"))
    if not os.path.exists(osp.join(args.output_dir, "train2017")):
        os.makedirs(osp.join(args.output_dir, "train2017"))
    if not os.path.exists(osp.join(args.output_dir, "val2017")):
        os.makedirs(osp.join(args.output_dir, "val2017"))

    # 获取目录下所有的.jpg文件列表
    feature_files = glob.glob(osp.join(args.input_dir, "*.jpg"))
    print('| Image number: ', len(feature_files))

    # 获取目录下所有的joson文件列表
    label_files = glob.glob(osp.join(args.input_dir, "*.json"))
    print('| Json number: ', len(label_files))
    

    # feature_files:待划分的样本特征集合    label_files:待划分的样本标签集合    test_size:测试集所占比例 
    # x_train:划分出的训练集特征      x_test:划分出的测试集特征     y_train:划分出的训练集标签    y_test:划分出的测试集标签
    x_train, x_test, y_train, y_test = train_test_split(feature_files, label_files, test_size=0.3)
    print("| Train number:", len(y_train), '\t Value number:', len(y_test))

    # 把训练集标签转化为COCO的格式，并将标签对应的图片保存到目录 /train2017/
    print("—"*50) 
    print("| Train images:")
    to_coco(args,y_train,train=True)
    
    # 把测试集标签转化为COCO的格式，并将标签对应的图片保存到目录 /val2017/ 
    print("—"*50)
    print("| Test images:")
    to_coco(args,y_test,train=False)
    

if __name__ == "__main__":
    print("—"*50)
    main()
    print("—"*50)

 文件结构如图：

 对数据集进行转换：

python labelme2coco.py --input_dir images --output_dir coco --labels labels.txt

该脚本会在执行过程中，自动会训练集和数据集进行划分，默认比例7:3，可自己修改test_size=0.3值进行调整。

执行结束后，会自动生成coco文件夹，进入文件夹可看到如下图的目录结构：

 四、搭建yolox环境

创建conda环境：

conda create -n yolox python=3.8

进入conda环境：

conda activate yolox

安装pytorch：

在pytorch官网找到适合你的pytorch版本，官方地址：Previous PyTorch Versions | PyTorch

我的环境是cuda10.2、cudnn8.2，所以执行命令：

conda install pytorch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1 cudatoolkit=10.2 -c pytorch

下载yolox代码

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

安装yolox环境

pip install -r requirements.txt
pip install -v -e .  # or  python3 setup.py develop

下载预训练模型

在yolox的官方github地址中，下载一个官方提供的预训练模型，我下载的是YOLOX-l，然后将下载的yolox_l.pth模型文件放在premodel文件夹下，若没有该文件夹，则新建一个。

测试一下

python tools/demo.py image -f exps/default/yolox_l.py -c ./premodel/yolox_l.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

根据提示，在 ./YOLOX_outputs/yolox_l/vis_res/2022_09_01_16_07_29/ 目录下查看识别结果。

通过以上步骤，说明我们的YOLOX一切正常，接下来就可以训练我们的自定义数据了。

五、YOLOX训练自定义模型

第三步的时候生成了coco文件夹，大家还记着吧？将coco文件夹放在yolox的datasets目录下。

然后复制一份exps\example\custom\yolox_s.py命名为yolox_l.py，打开后将self.depth和self.width改为1.0，将self.num_classes改为1。

注意：为什么改成1.0？大家可以参考看下exps\default\yolox_l.py文件，yolox_l和yolox_s最大的区别也就是网络深度不同。

我也看到过不少人说，self.num_classes=分类数+1，加1的原因是因为有__background__。大家只要按照我上面的脚本对数据集进行转换就行了，不存在+1的问题，标记了几个分类，self.num_classes就等于几。

顺便提供一下我的yolox_l.py配置文件内容：

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

from yolox.exp import Exp as MyExp


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 1.0
        self.width = 1.0
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # Define yourself dataset path
        self.data_dir = "datasets/coco"
        self.train_ann = "instances_train2017.json"
        self.val_ann = "instances_val2017.json"

        self.num_classes = 1

        self.max_epoch = 300
        self.data_num_workers = 4
        self.eval_interval = 1

一切就绪，开始训练：

python tools/train.py -f exps/example/custom/yolox_l.py -d 1 -b 4 --fp16 -o -c premodel/yolox_l.pth

-d参数为gpu数量

-b为batch-size，若训练过程出现cuda out of memeory，就适当调低。

训练时可能会报错：

错误1：AttributeError: module 'yolox.layers.fast_cocoeval' has no attribute 'InstanceAnnotation'

解决方法：
将文件E:\pythonFiles\YOLOX\yolox\evaluators\coco_evaluator.py中270行左右的

try:
    from yolox.layers import COCOeval_opt as COCOeval
except ImportError:
    from pycocotools.cocoeval import COCOeval
    logger.warning("Use standard COCOeval.")

更改为：

try:
    # from yolox.layers import COCOeval_opt as COCOeval
    from pycocotools.cocoeval import COCOeval
except ImportError:
    from pycocotools.cocoeval import COCOeval
    logger.warning("Use standard COCOeval.")

错误2：OSError: [WinError 1455] 页面文件太小,无法完成操作。

 

 再点击工程所在的盘——自定义大小（depend on 你的磁盘有多少空余，可以直接写满剩余空间。因为只是跑的时候占用虚拟内存，所以实际不会影响，放心）——设置（一定要点，不然没有效果！！）——确定

重启即可！！！

六、训练结果测试

训练结束后，可在YOLOX_outputs\yolox_l找到训练好的模型。

 可以加载自己的模型，测试一下识别效果。

python tools/demo.py image -f exps/example/custom/yolox_l.py -c ./YOLOX_outputs/yolox_l/latest_ckpt.pth --path ./datasets/coco/val2017/200042.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

 测试发现，识别一张图片大约需要0.9秒左右，接下来就可以搭建一下tensorrt的环境，来提高模型的推理速度。

七、安装TensorRT

什么是TensorRT

一般的深度学习项目，训练时为了加快速度，会使用多 GPU 分布式训练。但在部署推理时，为了降低成本，往往使用单个 GPU 机器甚至嵌入式平台（比如 NVIDIA Jetson）进行部署，部署端也要有与训练时相同的深度学习环境，如 caffe，TensorFlow 等。由于训练的网络模型可能会很大（比如，inception，resnet 等），参数很多，而且部署端的机器性能存在差异，就会导致推理速度慢，延迟高。这对于那些高实时性的应用场合是致命的，比如自动驾驶要求实时目标检测，目标追踪等。所以为了提高部署推理的速度，出现了很多轻量级神经网络，比如 squeezenet，mobilenet，shufflenet 等。基本做法都是基于现有的经典模型提出一种新的模型结构，然后用这些改造过的模型重新训练，再重新部署。TensorRT 则是对训练好的模型进行优化。 

7.1确认安装环境

使用以下命令查看你的cuda和cudnn的版本：

import torch
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())

 我电脑的环境是cuda10.2、cudnn8.2。

7.2下载安装包

下载地址：https://developer.nvidia.com/nvidia-tensorrt-download

我选择了TensorRT8版本，EA是提前发布的不稳定版本,GA是经过完备测试的稳定版。

找到windows版本的针对CUDA10.2的包，然后下载。 

下载完成后解压到D盘，如下：

 将以下地址配置到PATH环境变量中：

D:\TensorRT-8.2.5.1\lib
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\lib\x64

复制TensorRT-8.2.5.1目录下的bin、lib、include文件夹，复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2目录下。

# 安装TensorRT
cd TensorRT-8.2.5.1/python
pip install tensorrt-8.2.5.1-cp38-none-win_amd64.whl
 
# 安装UFF,支持tensorflow模型转化
cd TensorRT-8.2.5.1/uff
pip install uff-0.6.9-py2.py3-none-any.whl
 
# 安装graphsurgeon，支持自定义结构
cd TensorRT-8.2.5.1/graphsurgeon
pip install graphsurgeon-0.4.5-py2.py3-none-any.whl

八、安装torch2trt

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts    
bash build_contrib.sh

运行 python 测试，导入模块不报错就表明安装正确

import tensorrt
import uff

在import uff时如果出现了ModuleNotFoundError: No module named 'tensorflow'错误，那就使用pip安装一下：

pip install tensorflow-gpu

九、转换模型，生成engine文件

进入yolox目录，执行

python tools/trt.py -f exps/example/custom/yolox_l.py -c ./YOLOX_outputs/yolox_l/latest_ckpt.pth

转换过程非常慢，耐心等待十几分钟

 出现Converted TensorRT model done提示意味着转换成功。

九、yolox使用TensorRT进行推理

转换成功后会在YOLOX_outputs\yolox_l目录下生成model_trt.engine、model_trt.pth两个文件，此时我们就可以使用tensorRT加速推理过程了。

我分别以普通模式和tensorRT加速的模式对同一张图片进行识别：

#普通识别
python tools/demo.py image -f exps/example/custom/yolox_l.py -c ./YOLOX_outputs/yolox_l/latest_ckpt.pth --path ./datasets/coco/val2017/200042.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

#tensorRT加速识别
#不需要通过-c指定模型，看demo.py就明白，执行参数有--trt，默认加载model_trt.pth文件
python tools/demo.py image -f exps/example/custom/yolox_l.py --path ./datasets/coco/val2017/200042.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu --trt

识别结果：

 

 tensorRT的速度不是一般快啊，兄弟萌！！！

参考资料：

 windows下配置YOLOvX教程_拔牙的萌萌鼠的博客-CSDN博客_windows运行yolo

labelme转coco数据集 - gy77 - 博客园

YOLOX module 'yolox.layers.fast_cocoeval' 解决方案 - SegmentFault 思否

windows下配置YOLOvX教程_拔牙的萌萌鼠的博客-CSDN博客_windows运行yolo

YOLOV5训练模型报错：OSError: [WinError 1455] 页面文件太小,无法完成操作

TensorRT安装及使用教程_宗而研之的博客-CSDN博客_tensorrt安装

YOLOX Window10 TensorRT 全面部署教程_夜夜夜Best的博客-CSDN博客_tensorrt部署

 安装torch2trt 和TensorRT 踩坑记录_下大禹了的博客-CSDN博客_torch2trt安装

win10安装TensorRT - 夕西行 - 博客园 (cnblogs.com)