使用TensorRT加速Yolo

使用TensorRT加速Yolov5模型
使用版本为TensorRT-Yolov5-v7.0 , Yolov5 v7.0

以yolov5s.pt为例，进行目标检测

下载工程源码

# 按照版本进行项目克隆 yolo与tensorrt的版本需一致
git clone -b v7.0 https://github.com/ultralytics/yolov5.git
git clone -b yolov5-v7.0 https://github.com/wang-xinyu/tensorrtx.git

获取 wts 文件

# 由pt转换
cd yolov5/
wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt
cp [PATH-TO-TENSORRTX]/yolov5/gen_wts.py .
python gen_wts.py -w yolov5s.pt -o yolov5s.wts

也可以直接从model zoo中下载

编译tensorrt工程

cd [PATH-TO-TENSORRTX]/yolov5/
# 若为自己训练的模型需要在 yololayer.h 中进行响应更改类别数量等
mkdir build
cd build
cp [PATH-TO-ultralytics-yolov5]/yolov5s.wts . 
cmake ..
make

我在这 make 中出现了报错 /usr/bin/ld: warning: libnvcudla.so, needed by /usr/local/cuda/lib64/libcudla.so.1, not found (try using -rpath or -rpath-link)
这里我查出来是因为我的这个地址是存在软连接的，而不是实存在的路径，后我在 CMakeLists.txt 中把有关的cuda路换为不存在软连接的路径，问题解决

转换 wts 为 engine 文件

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

注意两个文件的路径问题

进行预测

在获取 engine 文件后就可以利用其进行预测，使用 engine 进行预测将自动调用tensorrt 如在之前的路径下

./yolov5_det -d yolov5.engine ../images

即可对它提供的两张照片_zidane.jpg 和 _bus.jpg进行预测

也可以只用python进行调用

# 需要下载python-tensorrt, pycuda, etc.
python yolov5_det_trt.py

# 或使用CUDA-python代替pycuda
python yolov5_det_trt_cuda_python.py

小捷径

如果不使用c++进行推理，而只用tensorrt对python项目进行加速，是可以直接在yolo项目中使用 export.py 将 pt 直接转为 engine，程序会直接调用tensorrt的包进行转换，具体使用例子如

python export.py --include engine --weights yolov5s.pt --device 0 --half

其中我指定了几个常用参数
--include engine 为要转换的目标模型，tensorrt的目标为engine
--weights yolov5s.pt 为被转换的模型文件，这里我们欲将yolov5s.pt进行转换
--device 0 为指定设备为cuda: 0，由于tensorrt限制，这里必须使用gpu进行export，可根据gpu设备号修改数字
--half 表示半精度推理
除了以上，还有非常多参数可以根据需求选择，具体可查看yolo项目 export.py 中的参数

TensorRT的使用

只需要在使用模型推理时，指定模型权重文件为对应的 .engine 即可

其他

如分类、分割等与其类似，可以参考官方介绍 进行操作