博客

1.下载必要的脚本

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

git clone -b archive https://github.com/ultralytics/yolov3.git

将以上两个脚本放置同一目录下；


或者采用笔者上传整理好的脚本文件



链接地址：


https://download.csdn.net/download/weixin_44533869/21026929

2.配置环境

1. 下载pytorch离线安装包，放置pkg目录下。
   
   torch-1.6.0-cp36-cp36m-linux_aarch64.whl
   
2. 下载torchvision源码库，放置pkg目录下。
   
git clone --branch v0.7.0 https://github.com/pytorch/vision torchvision

3. 执行安装脚本
   
bash install.sh


3.Yolo转TensorRT Engine流程

yolo -> wts(Pytorch Yolo) -> engine(TensorRT engine)

4. yolo转wts

进入yolov2wts目录

1. 将.cfg拷贝至cfg目录，将.weights拷贝至weights目录。（两个文件前缀必须相同，如yolov4_final.cfg,yolov4_final.weight，以下也以此为例)。
2. 执行生成yolov4_safe.wts命令。

python3 gen_wts.py yolov4_final

此处可以在本地执行，也可以在远端服务器执行

3. wts目录下将生成yolov4_safe.wts后缀的文件。

5. 下载Tensorrt

https://developer.nvidia.com/nvidia-tensorrt-7x-download


此处以Tensorrt-7x为例



笔者采用TAR格式的package包

6.安装Tensorrt

1. 解压TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz
2. 将下面环境变量写入环境变量文件~/.bashrc并保存

vim ~/.bashrc
export LD_LIBRARY_PATH=TensorRT解压路径/lib:$LD_LIBRARY_PATH
例如：
export LD_LIBRARY_PATH=/home/zhougen001/TensorRT-7.2.3.4/lib:$LD_LIBRARY_PATH

3. 使刚刚修改的环境变量文件生效

source ~/.bashrc

4. 复制TensorRT路径下/lib、/include文件夹到对应系统文件夹

sudo cp -r ./lib/* /usr/lib             # 在TensorRT-7.2.3.4路径下执行
sudo cp -r ./include/* /usr/include     # 在TensorRT-7.2.3.4路径下执行

5. 验证Tensorrt安装情况
   
   如果要使用python接口的tensorrt，则需要安装pycuda

pip install pycuda

测试Tensorrt

 $ python
   >>>import tensorrt
   >>>tensorrt.__version__

7. wts转TensorRT engine

进入wts2engine/yolov4目录

1. 新建build目录，并进入build目录。
2. 拷贝yolov4_safe.wts文件至build目录。

cp ../../../yolov2wts/wts/yolov4_safe.wts ./

3. 修改yololayer.h头文件

static constexpr int CLASS_NUM = 12;  //默认为80类别，修改为你自己训练的模型类别数量。yolov4_final为8类

cmake ..
make
生成TensorRT engine文件
sudo ./yolov4 -s  yolov4_safe

这里会生成yolov4_safe.engine文件。

4. 验证模型输出

sudo ./yolov4 -d yolov4_safe ../../../sample/

在build目录下生成推理结果