yolo5+tensorrt加速

yolo5+tensorrt加速

1.使用前注意

搭建环境需要下载tensorrt的包 yoloV5的版本和tensorrt的包要对应 本教程使用的是yolo5 V5.0的版本，所以tensorrt加速的包也要使用V5.0的 这是tensorrt下载链接：https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

2. 开始使用

获取tensorrtx(可能因为网络访问的原因会出错,需要多试几次)

git clone -b yolov5-v7.0 https://github.com/wang-xinyu/tensorrtx.git

把tensorrt/yolov5下的gen_wts.py复制到yolov5的文件夹下

执行gen_wts.py生成.wts文件。

python3 gen_wts.py -w yolov5s.pt -o yolov5s.wts

到目录tensorrtx下的yolov5文件夹，创建一个build文件,并进入

mkdir build

cd build

cmake ..

将yololayer.h里的CLASS_NUM修改成你的。官方用的是coco数据集，所以默认是80。(使用的是官方的可以忽略这步)

执行makeFile。（每次修改为CLASS_NUM都要make一次）

make -j2

从yoloV5的文件路径下的wts文件复制到tensorrtx/yolov5里。

生成.engine文件

sudo ./yolov5_det -s ../yolov5s.wts yolov5s.engine s

3.测试是否可以实现加速

方法1：进入tensorrtx/yolo5/build的文件夹下，运行命令

sudo ./yolov5_det -d yolov5s.engine ../images

可以和在yolov5文件夹下运行

python3 detect.py

进行一个对比，可以很清楚的看到加入了tensorrtx后，识别的图片时间大大减小

方法2：进入tensorrtx/yolo5的文件夹下，运行命令

python3 yolov5_trt.py

同样和可以和在yolov5文件夹下运行

python3 detect.py

进行一个对比，可以很清楚的看到加入了tensorrtx后，识别的图片时间大大减小