ubuntu18下TensorRT加速YOLOV5推理(超详细)

安装TensorRT

TensorRT是Nvidia官方给的C++推理加速工具,如同OpenVINO之于Intel。支持诸多的AI框架,如Tensorflow,Pytorch,Caffe,MXNet等。

可以通过以下链接进行TensorRT的下载安装,根据自己的cuda版本安装相对应的版本。

下载连接:https://developer.nvidia.com/tensorrt

下载这个需要注册nvidia账号,按要求注册就好了。
下载完成后解压:

tar -xvzf TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz

进行export:

export TRT_RELEASE=`pwd`/TensorRT-6.0.1.5
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TRT_RELEASE/lib

安装Python模块:

cd  TensorRT-6.0.1.5/python/
pip install tensorrt-6.0.1.5-cp37-none-linux_x86_64.whl

安装完毕!

部署YOLOV5-tensorrt

Yolov5使用TensorRT的过程,实际就是在其API下复现yolov5的过程。已经有大佬实现了并且在github上开源了项目,大家可以直接下载使用了。

项目地址:https://github.com/wang-xinyu/tensorrtx

(1)下载yolov5和 yolov5 tensorrt:

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git

(2)将tensorrt/yolov5/gen_wts.py拷贝至yolov5下:

sudo cp -r tenorrtx/yolov5/gen_wts.py {路径}yolov5

(3)生成yolov5s.pt对应的yolov5x.wts:

cd yolov5
python gen_wts.py -w yolov5ws.pt -o yolov5s.wts

(4)修改tensorrtx中yolov5下的CMakeLists.txt:

include_directories("/home/tensorrt_tar/TensorRT-6.0.1.5/include")    #具体路径根据实际安装路径来写
link_directories(/home/tensorrt_tar/TensorRT-6.0.1.5/lib)

5)以原始yolov5s模型为例运行:
C++接口:

cd tensorrtx/yolov5/
mkdir build && cd build
sudo cp -r yolov5/yolov5s.wts build
cmake ..
make

//由wts生成engine文件
./yolov5 -s yolov5s.wts yolov5s.engine s
//利用engine文件进行推理,此处输入图片所在文件夹,即可得到预测输出
./yolov5 -d yolov5s.engine ../samples

Python接口:

python yolov5_trt.py

说明:
得到engine文件后即可进行推理,但默认输入是图片文件夹,而一般来说实时性项目多用摄像头做输入源,故在yolov5.cpp中魔改一下再重新编译就行。此外IRuntime, ICudaEngine, IExecutionContext等核心类源码中每次有图片输入都会创建一次,严重拖慢速度,改成仅初始化阶段创建,后续再有图片输入时保持可以提高推理速度。

你可能感兴趣的:(笔记,深度学习,caffe,神经网络,图像识别,目标检测)