Jetson Nano使用tensorrtx部署YOLOV5

1、为nano内置的cuda配置环境变量

• nano内置好了cuda,但需要配置环境变量才能使用，打开终端添加环境变量即可。
  添加前查看是否是cuda10.2，如果是按照下面的示例填写即可，如果不是更换下面示例中cuda版本即可。

 打开.bashrc文件
vim ~/.bashrc

在文件末尾添加：
export PATH=/usr/local/cuda-10.2/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export CUDA_ROOT=/usr/local/cuda

更新
source ~/.bashrc

• 配置完成后使用nvcc -V命令查看是否配置成功，如果出现以下字样则说明配置成功。

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Feb_28_22:34:44_PST_2021
Cuda compilation tools, release 10.2, V10.2.300
Build cuda_10.2_r440.TC440_70.29663091_0

2、安装所需要用到的库和环境依赖

sudo apt-get install build-essential make cmake cmake-curses-gui
sudo apt-get install git g++ pkg-config curl
sudo apt-get install libatlas-base-dev gfortran libcanberra-gtk-module libcanberra-gtk3-module
sudo apt-get install libhdf5-serial-dev hdf5-tools 
sudo apt-get install nano locate screen

sudo apt-get install libfreetype6-dev 
sudo apt-get install protobuf-compiler libprotobuf-dev openssl
sudo apt-get install libssl-dev libcurl4-openssl-dev
sudo apt-get install cython3

sudo apt-get install libavcodec-dev libavformat-dev libswscale-dev
sudo apt-get install libxvidcore-dev libavresample-dev
sudo apt-get install libtiff-dev libjpeg-dev libpng-dev

• 环境依赖配置好以后安装pycuda

pip3 install pycuda

3、修改Nano显存

• 打开终端输入

sudo gedit /etc/systemd/nvzramconfig.sh

• 修改nvzramconfig.sh文件
  将

mem = $((("${totalmem}"/2/"${NRDEVICES}")*1024))

修改为

mem = $((("${totalmem}"*2/"${NRDEVICES}")*1024))

4、获取tensorrt开源库

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

在终端输入以下命令获取tensorrtx
git clone https://github.com/wang-xinyu/tensorrtx.git

• 注意：下载tensorrtx的时候一定要下载与自己使用的yolov5版本相对应的tensorrtx版本

5、编译tensorrtx

1. 进入到tensorrtx目录

cd ~/tensorrtx/yolov5/
如果是手动下载的则使用：
cd ~/tensorrtx-master/yolov5/

2. 根据自己的情况修改yololayer.h和yolov5.cpp

• 例如在yololayer.h修改检测类别数和输入图像尺寸

static constexpr int CLASS_NUM = 80;   // 类别种类
static constexpr int INPUT_H = 640;    // 输入图像的高
static constexpr int INPUT_W = 640;   // 输入图像的宽

3. 创建build目录并进入

mkdir build
cd build

4. 编译

cmake ..

make -j4

6、生成engine文件并进行测试

1. 生成.wts文件

• 将/tensorrtx/yolov5/中的gen_wts.py复制到yolov5项目的根目录下

在终端中执行：
python3 gen_wts.py -w yolov5s.pt -o yolov5s.wts
-w：训练好的yolov5模型的路径
-o：输出的 .wts 文件路径

• 注：这一步可以在电脑上进行，也可以在Nano上进行，个人推荐在电脑上进行。

2. 生成engine文件

• 将上一步生成的.wts文件放入/tensorrtx/yolov5/build/目录下
• 在终端中执行

sudo ./yolov5 -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]

以yolov5s模型为例：
sudo ./yolov5 -s yolov5s.wts yolov5s.engine s

3. 测试

• 在终端中执行

sudo ./yolov5 -d [.engine] [image folder]

例：
sudo ./yolov5 -d yolov5s.engine ../samples/
samples：图片的存放路径

7、如何通过Python调用

• 在/tensorrtx/yolov5/中tensorrtx作者给了两个python版本的示例分别是：yolov5_trt.py和yolov5_trt_cuda_python.py。