当前位置：首页 > news >正文

yolov5 自训练模型转 tensorrt 及测试

news 2025/6/9 16:03:24

文章目录

- 一、准备yolov5自训练模型
- 二、下载tensorrtx源码
- 三、从pt模型生成wts模型
- 四、修改tensorrtx源码
- - 4.1 修改 yololayer.h
  - 4.2 修改 yolov5.cpp
- 五、编译运行tensorrt/yolov5

在yolov5训练完模型后，将其转换为tensorrt格式，可以加速推理时间，相关转换及测试步骤如下

一、准备yolov5自训练模型

准备一个自己训练好的yolov5模型，如：

YOLOv5\runs\train\exp_ex_bu_77epoch_x\weights\best.pt

二、下载tensorrtx源码

wang-xinyu/tensorrtx
下载完成后，进入 tensorrt/yolov5 目录
在这里插入图片描述

三、从pt模型生成wts模型

先把 tensorrt/yolov5 目录的 gen_wts.py 文件复制到 yolov5 主目录，然后执行该脚本进行转换，将生成一个wts格式的模型

cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5
cd {ultralytics}/yolov5
python gen_wts.py -w runs/train/exp_ex_bu_77epoch_x/weights/best.pt -o yolov5x.wts
// a file 'yolov5x.wts' will be generated.

四、修改tensorrtx源码

进入 tensorrt/yolov5 目录

4.1 修改 yololayer.h

根据实际情况，修改20-22行的类别数目及网络输入尺寸
在这里插入图片描述

4.2 修改 yolov5.cpp

根据实际情况，修改11-16行的量化类型，GPU id号，NMS阈值，置信度阈值等参数
在这里插入图片描述

五、编译运行tensorrt/yolov5

进入 tensorrt/yolov5 目录，需要在 CMakeLists.txt 中设置TensorRT 正确路径。在服务器 ~/.bashrc 中设置的tensorrt 路径是 /usr/local/lib/TensorRT-7.0.0.11/lib，但编译时会出错，提示找不到 NvInfer.h 文件，解决方法是先通过locate全局查找是否存在 NvInfer.h 文件，通过以下命令查看，

locate NvInfer.h

在这里插入图片描述

可以看到该文件在另一个目录，说明此处的tensorrt才是完整的，
因此在CMakeLists.txt中，将TensorRT路径改成这里的

/home/xxx/xxx/YoloV5TernsorRTDemo/TensorRT-7.0.0.11/include/
/home/xxx/xxx/YoloV5TernsorRTDemo/TensorRT-7.0.0.11/lib/

注意，上述路径的xxx要根据自己的情况进行修改。
![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/2ccabafe76b74c38b39230b965783ea9.png

接下来重新进行编译

cd {tensorrtx}/yolov5
mkdir build
cd build
cmake ..
make

执行完后将在build目录生成yolov5可执行文件
然后利用这个可执行文件可以将第三步的wts文件转换成engine文件，并进行推理测试

# 复制wts到当前目录
cp {ultralytics}/yolov5/yolov5x.wts {tensorrtx}/yolov5/build
# 生成engine模型，最后一个x代表yolov5模型尺寸，根据实际情况修改
sudo ./yolov5 -s yolov5x.wts yolov5x.engine x
# engine模型推理测试
sudo ./yolov5 -d yolov5x.engine ./test_imgaes