文章目录
- 一、准备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要根据自己的情况进行修改。
接下来重新进行编译
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