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_splittry: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 -1class_name = line.strip()   # strip() 方法用于移除字符串头尾指定的字符（默认为空格或换行符）或字符序列。if class_id == -1:assert class_name == "__ignore__"   # background:0, class1:1, ,,continueclass_name_to_id[class_name] = class_iddata["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 areasegmentations = collections.defaultdict(list)  # for segmentationfor 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] | maskelse:masks[instance] = maskif shape_type == "rectangle":(x1, y1), (x2, y2) = pointsx1, 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 = instanceif cls_name not in class_name_to_id:continuecls_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 osfrom yolox.exp import Exp as MyExpclass Exp(MyExp):def __init__(self):super(Exp, self).__init__()self.depth = 1.0self.width = 1.0self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]# Define yourself dataset pathself.data_dir = "datasets/coco"self.train_ann = "instances_train2017.json"self.val_ann = "instances_val2017.json"self.num_classes = 1self.max_epoch = 300self.data_num_workers = 4self.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 COCOevallogger.warning("Use standard COCOeval.")

更改为：

try:# from yolox.layers import COCOeval_opt as COCOevalfrom pycocotools.cocoeval import COCOeval
except ImportError:from pycocotools.cocoeval import COCOevallogger.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)