代码:
https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5(v3.0版)
https://github.com/ultralytics/yolov5(v3.0版)
环境:cuda:10.2 ,cudnn:8.0.4, tensorrt :7.1.3,pytorch:1.7.0
一、生成转tensorrt的中间文件 yolov5.wts
yolov5s.pt (v3.0)放入weights文件夹
在yolov5-3.0工程中运行:
python gen_wts.py
生成wts文件
二、编译yolov5并生成tensorrt模型yolov5.engine
1、CMakeList.txt中对tensorrt路径进行指定,不然会报错找不到nvinfer
include_directories(/home/yuxin/TensorRT-7.1.3.4/include)
link_directories(/home/yuxin/TensorRT-7.1.3.4/lib)
2、./tensorrt/yolov5/下新建个samples文件夹,把需要测试的图片放进去
3、yolov5s.wts复制到 tensorrtx-yolov5-v3.0/yolov5
在终端下运行:
mkdir build
cd build
cmake ..
make
sudo ./yolov5 -s // serialize model to plan file i.e. 'yolov5s.engine'
sudo ./yolov5 -d ../samples // deserialize plan file and run inference, the images in samples will be processed.