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

参考代码
代码支持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. 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 !

生成wts文件

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

cp  gen_wts.py /home/yolov5/

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

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

编译engine文件

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

mkdir build
cd build
cp /home/yolov5/yolov5s_best.wts /home/yolov5/build
cmake ..
make

根据自己使用的模型类别运行生成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

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

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

你可能感兴趣的:(目标检测,深度学习,目标检测)