【Ubuntu】yolov5-6.0+tensorRT在Ubuntu(x86_64)部署推理

1 yolov5-6.0虚拟环境搭建

这部分的内容可以看笔者的博客，链接🔗：【YOLOv5】6.0环境搭建（不定时更新）

2 安装OpenCV

这部分的内容可以看笔者的博客，链接🔗：【Ubuntu】Ubuntu18.04安装OpenCV4.6.0（x86_64 / Jetson Xavier NX arm64）

3 安装TensorRT

查看Ubuntu内核的安装或者说CPU架构，在终端输入

arch
# 或者
uname -a
1
2
3

1. aarch64：ARM 架构
2. x86_64: X86 架构

笔者以x86_64: X86 架构为例实现安装（关于Jetson Xavier NX arm64的部署，笔者测试过与x86_64: X86 架构基本一致）

查看自己的CUDA版本

nvcc -V
1

官网上下载对应的TensorRT
https://developer.nvidia.com/tensorrt
官网链接🔗：https://developer.nvidia.com/nvidia-tensorrt-7x-download

解压并安装

# 提取到此处
tar -xzvf TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.1
# 打开环境变量文件
vim ~/.bashrc 
# 将下面三个环境变量写入环境变量文件并保存
export LD_LIBRARY_PATH=/home/ll/TensorRT-7.2.3.4/lib:$LD_LIBRARY_PATH
export CUDA_INSTALL_DIR=/usr/local/cuda-10.2
export CUDNN_INSTALL_DIR=/usr/local/cuda-10.2
# 使刚刚修改的环境变量文件生效
source ~/.bashrc
# 进入虚拟环境（其中yolov5为自定义的虚拟环境的名称）
conda activate yolov5
python  # 查看python的版本
quit()  # 退出python
# 进入TensorRT-7.2.3.4/python，python选择与虚拟环境一致的版本，如下图所示
cd TensorRT-7.2.3.4/python
pip install tensorrt-7.2.3.4-cp38-none-linux_aarch64.whl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

#  查看tensorrt是否可以正常导入
python
import tensorrt
tensorrt.__version__
quit()
1
2
3
4
5

成功安装TensorTR=7.2.3.4

4 实现推理部署（yolov5-6.0+tensorrt）

编程软件笔者推荐vscode，vscode安装可以看笔者的博客，链接🔗：【Ubuntu】安装Anaconda+vscode
（1）文件准备

1. tensorrt/yolov5-6.0源码
2. yolov5-6.0源码
3. yolov5s.pt 预训练模型（为了测试，后续可以使用自训练的模型）
   下载预训练模型yolov5s.pt，将yolov5s.pt文件放到yolov5-6.0/weights文件

#  模型下载链接
https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
1
2

（2）yolov5s.pt–>yolov5s.wts
将tenorrtx/yolov5/gen_wts.py复制到yolov5-6.0文件中，并运行下面的代码，生成yolov5s.wts

code ./yolov5-6.0  #在终端输入，使用 vscode打开源码项目（根据自己的路径修改）
1

conda activate yolov5  # 激活环境
python gen_wts.py --weight weights/yolov5s.pt  # 在weights文件夹下生成yolov5s.wts
1
2

（3）TensorTR编译，生成执行文件yolov5

code ./tensorrtx-yolov5-v6.0/yolov5  # 在终端输入，使用 vscode打开源码项目（根据自己的路径修改）
mkdir build
cd build  
1
2
3

修改CMakeLists.txt中cuda和tensorrt对应的路径，否则会编译报错

开始编译，终端出现如图所示的信息，则成功编译

cmake ..
make
1
2

执行完在build文件夹下生成yolov5执行文件，如图所示，并将（2）生成的yolov5s.wts文件复制到该build文件夹下

（4）生成引擎文件yolov5s.wts–>yolov5s.engine

sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
1

执行完后在build文件夹下生成yolov5s.engine引擎文件

（5）部署后推理测试

sudo ./yolov5 -d yolov5s.engine ../samples
1

至此，推理部署完毕，自训练模型按照上述操作即可实现推理加速

（6）在python中加载并运行tensorrt模型

#  安装pycuda
pip install pycuda -i https://pypi.tuna.tsinghua.edu.cn/simple
# 当前目录为  ./tensorrtx-yolov5-v6.0/yolov5
python yolov5_trt.py
1
2
3
4