linux下pythorch训练yoloV5模型转换为tensorRT模型_将模型转换linux

参考代码
代码支持yolov3-yolov5，yolov5支持v1.0-v6.0。本文以yolov5的v6.0进行测试。

yoloV5模型准备

首先需要配置yoloV5运行环境，具体可以参考如下链接，配置完成后，假如克隆的代码在/home/yolov5文件夹下，进入文件夹训练自己的模型，得到yolov5s_best.pt模型文件。

模型转换代码准备

从如下链接克隆tensorrtx代码，里面包含多种模型的转换，克隆的代码放在/home/tensorrtx/文件夹下。

参数配置

/home/tensorrtx/目录中，需要根据自己训练模型修改如下配置

1. yololayer.h中定义输入形状、类别数、最大box数等，具体参数如下：

    static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000;
    static constexpr int CLASS_NUM = 80;
    static constexpr int INPUT_H = 640;  // yolov5's input height and width must be divisible by 32.
    static constexpr int INPUT_W = 640;
1
2
3
4

2. yolov5.cpp修改

//选择INT8/FP16/FP32 
#define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
//修改NMS阈值
#define NMS_THRESH 0.4
//修改置信度
#define CONF_THRESH 0.5
//修改BATCH_SIZE 
#define BATCH_SIZE 1
//修改输入图像最大尺寸，超过该尺寸会越界
#define MAX_IMAGE_INPUT_SIZE_THRESH 3000 * 3000 // ensure it exceed the maximum size in the input images !
1
2
3
4
5
6
7
8
9
10

生成wts文件

进入/home/tensorrtx/目录，将gen_wts.py文件拷贝到/home/yolov5文件夹下。

cp  gen_wts.py /home/yolov5/
1

conda activate yolov5 激活yolov5环境，进入/home/yolov5文件夹，运行如下命令得到wts文件。

python gen_wts.py -w yolov5s_best.pt -o yolov5s_best.wts
1

编译engine文件

进入/home/tensorrtx/目录，编译后生成yolov5可执行文件

mkdir build
cd build
cp /home/yolov5/yolov5s_best.wts /home/yolov5/build
cmake ..
make
1
2
3
4
5

根据自己使用的模型类别运行生成engine文件，使用命令sudo ./yolov5 -s [.wts] [.engine] s，s可以根据自己的情况换为n/s/m/l/x/n6/s6/m6/l6/x6等。

sudo ./yolov5 -s yolov5s_best.wts yolov5s_best.engine s
1

然后可以运行sudo ./yolov5 -d [.engine] [image folder]进行模型验证

sudo ./yolov5 -d yolov5s_best.engine ../samples
1