yolov4模型经tensorrt加速

最新推荐文章于 2024-04-22 17:25:05 发布

沐雲小哥

最新推荐文章于 2024-04-22 17:25:05 发布

阅读量1.3k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： Yolo

本文链接：https://blog.csdn.net/weixin_44533869/article/details/119617858

Yolo 专栏收录该内容

15 篇文章

订阅专栏

该博客介绍了如何将YOLOv4模型转换为TensorRT engine，以便于在CUDA环境下进行高性能推理。首先，从github下载必要的脚本并放置在同一目录下，然后配置PyTorch和TensorRT的环境。接着，通过YOLO模型的.cfg和.weights文件生成.wts权重文件。下载并安装TensorRT后，修改相关头文件并编译，最后生成TensorRT engine文件。通过运行测试脚本验证模型输出，完成模型的优化过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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.配置环境

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

3.Yolo转TensorRT Engine流程

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

4. yolo转wts

进入yolov2wts目录

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

python3 gen_wts.py yolov4_final

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

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

5. 下载Tensorrt

https://developer.nvidia.com/nvidia-tensorrt-7x-download
此处以Tensorrt-7x为例
在这里插入图片描述
笔者采用TAR格式的package包

6.安装Tensorrt

解压TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz
将下面环境变量写入环境变量文件~/.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

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

source ~/.bashrc

复制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路径下执行

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

pip install pycuda

测试Tensorrt

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

7. wts转TensorRT engine

进入wts2engine/yolov4目录

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

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

修改yololayer.h头文件

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

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

这里会生成yolov4_safe.engine文件。
在这里插入图片描述

验证模型输出

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

在build目录下生成推理结果
在这里插入图片描述