yolov5模型部署：Nvidia使用TensorRT部署yolov5s模型

点上方计算机视觉联盟获取更多干货

仅作学术分享，不代表本公众号立场，侵权联系删除

转载于：作者丨梁彦哲@知乎（已授权）

来源丨https://zhuanlan.zhihu.com/p/386881025

编辑丨极市平台

AI博士笔记系列推荐

周志华《机器学习》手推笔记正式开源！可打印版本附pdf下载链接

一、TensorRT是什么？

TensorRT 是由 Nvidia 推出的 GPU 推理引擎(GIE: GPU Inference Engine)。和通用的深度学习框架不同，TensorRT 只提供前向传播，即推理的功能，而没有训练的功能。实际上，训练的场景下，资源通常不会特别紧张，因为是离线的操作(不要求实时反馈)，即使出现了一时的资源不足也可以通过增加计算时长来弥补，所以训练框架的资源消耗一般不会有很强的优化。而网络的部署通常对资源更加敏感，算力与内存都是需要考虑的因素。TensorRT 是一个旨在极致优化 GPU 资源使用的深度学习推理计算框架。其工作主要分为两个阶段：建造阶段(build phase)和执行阶段(compile phase)。

在建造阶段，TensorRT 接收外部提供的网络定义(也可包含权值 weights)和超参数，根据当前编译的设备进行网络运行的优化(optimization), 并生成推理引擎 inference engine(可以以 PLAN 形式存在在硬盘上)；

在执行阶段，通过运行推理引擎调用 GPU 计算资源——整个流程如下所示:

二、为什么要使用TensorRT

1.满足实时性要求

使用TensorRT加速推理是嵌入式AI非常重要的一环。具体来说以我自己1660Ti笔记本跑模型尚且不能达到满足要求的帧率，更别指望TX2有多好的效果。经实际测试，TX2端TensorRT加速后可至40帧左右，满足实时性要求。

2.TX2环境限制（关键！）

以最新的TX2刷机包Jetpack4.5.1为例，其封装好的CUDA版本为10.0，而我使用的yolov5 4.0版本需CUDA11.0，且yolov5所依赖的某些依赖尚无支持arm架构的版本，在TX2上直接跑yolov5是根本不可能的！必须使用TensorRT API对网络进行“复现”。

三、如何使用TensorRT

1.前期准备与说明

使用TensorRT的过程，实际就是在其API下复现yolov5的过程。对于任意一个已训练好的神经网络，我们需要得知其网络结构(backbone, neck等)与训练权重。训练权重已由系列教程的第二篇中得到，网络结构的获取方法主要有以下几种：1.使用TF-TRT，将TensorRT集成在TensorFlow中 2.使用ONNX2TensorRT，即ONNX转换trt的工具 3.手动构造模型结构，然后手动将权重信息挪过去，非常灵活但是时间成本略高，有大佬已经尝试过了：tensorrtx 其中前两种常常会有遇到不支持的结构/层的情况，解决该问题非常费时。而使用TensorRT API手动构造模型结构是最为稳妥、对模型还原度最高、精度损失最少的一种方法。万幸有dalao已将此开源，在此列出项目链接并表示感谢！

wang-xinyu/tensorrtxgithub.com

可以看到该项目除支持yolov5外，resnet, yolov3, yolov4等均支持，如果使用其他网络结构的同学也可以加以利用。

2. 使用流程

下载好该项目的项目的源码后，我们仅需要用到其中yolov5的子文件夹，其他为无关项。下载过程中要注意所下载的tensorrtx版本应与所选用的yolov5版本相匹配。

2.1 权重转换

打开gen_wts.py文件，修改对应权重路径

此步将会得到wts格式的yolov5权重，以供下一步生成engine文件

2.2 生成TensorRT Engine

// put yolov5s.wts into tensorrtx/yolov5// go to tensorrtx/yolov5// update CLASS_NUM in yololayer.h if your model is trained on custom dataset//注意在yololayer.h中改类别数！！
mkdir buildcd buildcmake ..make //编译完成
//由wts生成engine文件sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x or c gd gw]  // serialize model to plan file
//利用engine文件进行推理，此处输入图片所在文件夹，即可得到预测输出sudo ./yolov5 -d [.engine] [image folder]  // deserialize and run inference, the images in [image folder] will be processed.
// For example yolov5ssudo ./yolov5 -s yolov5s.wts yolov5s.engine ssudo ./yolov5 -d yolov5s.engine ../samples// For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yamlsudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25sudo ./yolov5 -d yolov5.engine ../samples

四、效果评估 如下是我在TX2端的实测视频

五、后续说明

1.得到engine文件后即可进行推理，但默认输入是图片文件夹，而一般来说实时性项目多用摄像头做输入源，故在yolov5.cpp中魔改一下再重新编译就行。此外IRuntime, ICudaEngine, IExecutionContext等核心类源码中每次有图片输入都会创建一次，严重拖慢速度，改成仅初始化阶段创建，后续再有图片输入时保持可以提高推理速度。

2.以后可能会融合deepsort做一下目标跟踪，敬请期待~

-------------------

END

--------------------

我是王博Kings，985AI博士，华为云专家、CSDN博客专家（人工智能领域优质作者）。单个AI开源项目现在已经获得了2100+标星。现在在做AI相关内容，欢迎一起交流学习、生活各方面的问题，一起加油进步！

我们微信交流群涵盖以下方向（但并不局限于以下内容）：人工智能，计算机视觉，自然语言处理，目标检测，语义分割，自动驾驶，GAN，强化学习，SLAM，人脸检测，最新算法，最新论文，OpenCV，TensorFlow，PyTorch，开源框架，学习方法...

这是我的私人微信，位置有限，一起进步！

王博的公众号，欢迎关注，干货多多

王博Kings的系列手推笔记（附高清PDF下载）：

博士笔记 | 周志华《机器学习》手推笔记第一章思维导图

博士笔记 | 周志华《机器学习》手推笔记第二章“模型评估与选择”

博士笔记 | 周志华《机器学习》手推笔记第三章“线性模型”

博士笔记 | 周志华《机器学习》手推笔记第四章“决策树”

博士笔记 | 周志华《机器学习》手推笔记第五章“神经网络”

博士笔记 | 周志华《机器学习》手推笔记第六章支持向量机（上）

博士笔记 | 周志华《机器学习》手推笔记第六章支持向量机（下）

博士笔记 | 周志华《机器学习》手推笔记第七章贝叶斯分类（上）

博士笔记 | 周志华《机器学习》手推笔记第七章贝叶斯分类（下）

博士笔记 | 周志华《机器学习》手推笔记第八章集成学习（上）

博士笔记 | 周志华《机器学习》手推笔记第八章集成学习（下）

博士笔记 | 周志华《机器学习》手推笔记第九章聚类

博士笔记 | 周志华《机器学习》手推笔记第十章降维与度量学习

博士笔记 | 周志华《机器学习》手推笔记第十一章稀疏学习

博士笔记 | 周志华《机器学习》手推笔记第十二章计算学习理论

博士笔记 | 周志华《机器学习》手推笔记第十三章半监督学习

博士笔记 | 周志华《机器学习》手推笔记第十四章概率图模型

点分享

点收藏

点点赞

点在看