项目经验掏心窝

接触不到真实的工业项目？

理论知识储备极佳，面对项目实践却无从下手？

工业项目部署落地流程完全大脑空白？

极市平台推出了新的内容专栏【项目经验掏心窝】，来听听小伙伴们的项目实践经验分享，他们在如何项目中打怪升级，又有哪些好用的实操tricks！

极市也非常欢迎小伙伴们踊跃自荐，来分享自己的项目经历、技术总结或者是比赛的获奖方案，稿件发布会有相应的稿酬奖励！

有意向的小伙伴欢迎添加运营的账号（ID：fengcall19）

作者丨DearPuff

编辑丨极市平台

一、引言

作为一名初入深度学习算法开发领域的玩家，“算法开发”这四个字似乎和我想象中的不太一样。作为Python的忠实拥护者，在我还没参与到这份工作之前，我想象中的深度学习算法开发其实只是针对当前的最新抑或是最强理论做出复现，用大白话来说就是：“我用Python能跑得出来就行了。”

但是真正进入工作之后，我发现这远远不一样，因为你永远也想不到你的深度学习算法到底会被用在什么地方。对于深度学习来说，算力永远还是第一要义，即便不是第一那也是第二，可是真正下放到工况之中，由于机器的成本控制等种种因素，开发出来的算法能有拥有一张英伟达显卡，俨然是上上签。更多数时候，算法还会被应用于边缘设备，这些设备往往都是非常低的算力，这些时候，就要有算法开发工作者站出来，为深度学习的部署贡献自己的力量。

就拿我之前经历过项目来说，我的目标设备是一张JetsonNano-DevKit，也就是大家平时在市面上见到最多的那种（图1）并且，由于任务需要，我需要在100毫秒级的时间内完成深度学习的推理工作。这时候，强大的算法开发能力显得尤为重要，因为如果单纯的把Python训练出来的模型放到上面去运行那是不大现实的，经过实测下来的算法运行时间、算力需求等或许会不符合任务要求。

因而，我们为何需要做算法开发，不仅是要满足任务的需求，更是要把设备的性能压榨到极致，并且在极致压榨的设备性能下还要有相对快的时间去完成你的任务，总结就是：你的算法需要减少算力需求、加快推理时间、保证算法精度。

（图1 JetsonNano图例，图源：英伟达官网）

二、常见的算法开发方法

讲完了我们为什么要做算法开发，接下来就该讲讲在当下开发中我们比较常用的几种算法开发方法。

1. 深度学习模型转ONNX

ONNX（Open Neural Network Exchange），是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch、MXNet）可以采用相同格式存储模型数据并交互。ONNX的规范及代码主要由微软，亚马逊，Facebook和IBM等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有：Caffe2，PyTorch，MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持ONNX。[1]

在当下，至少在我接触到工作里，基本上用的Pytorch模型，并且大部分我们看得到的技术帖中，深度学习模型转ONNX模型的方法也是基于Pytorch模型来写的。Pytorch模型在转成ONNX模型之后，在同硬件平台及同软件框架下，ONNX推理的优点在于推理速度会有小幅度提升，但缺点是会在模型加载阶段会耗费大量时间。所以，在Python环境下直接将Pytorch模型转为ONNX模型的方法适用于只需要一次加载模型的任务，比如长时间视频推理，超大量图片推理等，因为在模型一次加载之后可以持续进行加速推理。反之，对于只有小量图片，譬如在Python环境下，推理10张图片在ONNX模型上的推理时间只有100ms，但是模型本身加载需要600ms，这类任务就直接调用原有的Pytorch模型进行推理即可。

2. ONNX转TensorRT

在模型转为ONNX模型之后，我们仍然需要寻求一种更快更节省空间的推理方法，也就是基于TensorRT的方法。

TensorRT是英伟达（NVIDIA）开发的用于高性能深度学习推理的SDK，这个SDK可以为深度学习推理应用提供低延迟和高吞吐量。根据英伟达官网的描述，基于TensorRT的应用程序的执行速度可比CPU平台的速度快40倍。因此，得益于这个特性，这也是我们为什么要在算法开发中期望使用TensorRT方法来进行工作。

大多数情况下，我们都是先将模型转为ONNX模型之后再转为TensorRT模型进行推理。但某些比较成熟的算法，比如YOLOv5系列算法，在环境配置好之后，也可以通过Pytorch模型文件导出TensorRT模型文件。在英伟达官网的描述中，最主要的方法还是通过ONNX模型导出TensorRT模型。

3. 基于C++的TensorRT开发

其实TensorRT同时提供了C++和Python的API，并且这些API都具有相同的功能，但为什么着中强调基于C++的TensorRT开发。原因是Python API并非可以在所有平台上使用，并且英伟达官网也着重强调了C++ API可以更高效的满足某些需求。所以使用C++来进行TensorRT算法开发，也是当下主流的用来提升算法速度和减少计算消耗的一种方法。

不论是基于C++还是Python，TensorRT算法开发都遵循近乎相同的流程。那就是分为构建（Build）阶段和运行时（Runtime）阶段。从C++的角度出发，在构建和运行时阶段之间增加了一个反序列化（Deserializing）阶段。

以C++API举例，在构建阶段，首先要实例化ILogger接口，该接口用来接收所有警告信息，接着是使用ONNX解析器导入网络，生成engine/wts/trt文件也就是TensorRT模型，下一章中将会给出更简便的方法来构建TensorRT模型。在生成TensorRT模型后，得到engine文件（在下文中均用“引擎”代替“engine”），就需要构建引擎，至此整个构建阶段结束，进入反序列化阶段。

在反序列化阶段中，我们便可以创建运行时（Runtime）实例并且调用刚才创建的ILogger接口来接收所有错误信息，并且在这一步中，将模型读入缓冲区，也就是上文提到的运行时之前加入反序列化。

反序列化结束之后就可以使用初始化结束的模型进行推理，我们在这一步中加入cuda执行上下文（context）来进行推理任务，值得注意的是，一个引擎可以有多个上下文，意味着我们可以通过使用一组权重并多次创建上下文的行为来进行批量推理任务（当模型具有动态输入尺寸的时候不适用）。之后将处理好的数据调用cuda流并使用TensorRT进行推理，得到输出张量再进行后处理即可，当然数据的输入输出及cuda流的同步异步具体看任务需求，这里便不再赘述。

三、实战：基于YOLOv5-7.0 demo的算法开发演示

在本章中，将以基础演示的形式，在YOLOv5-7.0的demo基础上，来描述我们常用的几种算法开发手段。

1. 环境配置

工欲善其事必先利其器，为了保证在实践过程中尽量不踩坑，在本文中尽量采用目前最新的软件包来进行演示。

基本环境配置如下

（1）硬件配置：

CPU：i5-9600K 6C6T

GPU：RTX2060 6G

RAM：16G

（2）软件配置：

Windows11专业版 22H2

VisualStudio 2022 Community->安装工作负荷：“通用Windows平台开发”和“使用C++的桌面开发”

cuda 11.8

TensorRT 8.6.1.6

Anaconda2023.09

Python 3.10.13

torch 2.1.0+cu118

onnxruntime-gpu 1.16.0

具体如何配置，就不在这里赘述。不过值得注意的是，如果是想在Python环境下安装TensorRT，那就打开下载好的TensorRT安装包，把图2中箭头所指的四个文件夹中的对应自己Python版本的whl文件安装到自己的Python环境即可

（图2 TensorRT软件包示例）

2. Pytorch模型转ONNX

我们从GitHub中下载YOLOv5的最新源码、及yolov5s的预训练模型（图3），其实YOLOv5应该是更新到8.0的，但是不知什么原因回退到了7.0，不过对于实验是没有太多影响的。拿到源码后，我们根据requirements.txt中的内容安装好YOLOv5所需环境。

这里有一个小技巧，其实对于刚踏入深度学习的小伙伴来说，提前把VisualStudio装好能在Python环境的安装中少踩很多坑，原因是很多软件包都会依赖C环境编译及安装，不提前装好可能会有很多Python软件包安装不上，比如使用coco数据集训练YOLOv5时，就没法正确安装“pycocotools”这个软件包。当然，这也只是在Windows开发时比较容易遇到的坑，不过并不意味着Linux下不会出现类似的问题，总之在软件环境的安装下，两个平台的原理几乎都是相通的。

（图3 YOLOv5的releases界面，网址：）

由于我们拿到了预训练模型并且只是做一个简单的demo，这里便不再描述YOLOv5模型训练的过程，我们直接快进到将Pytorch模型转为ONNX模型这一步。

首先我们先将下载好的预训练模型文件yolov5s.pt，存放至我们想要的地方，在这里我将其存放至我源码目录下的/checkpoints目录（如图4）。

（图4 存放预训练模型）

接着我们打开YOLOv5源码中的export.py，移动至823行（如图5），在这里我们能看到关于export.py的所有参数，我们只需要更改这些参数便可对pt文件进行操作，将其转为其他模型文件。

（图5 export.py中的参数）

这么多参数或许会让我们眼花缭乱，由于我们是使用的预训练模型，且我们只需要导出ONNX模型，我们只需要更改825行的“weights”参数中的default值，将其改成我们上一步中存放的预训练模型的路径，在本文中是“ROOT / 'checkpoints/yolov5s.pt'”；以及849行中的default的值将其换成“['onnx']”//至于很多眼尖的伙伴会发现还可以更改device等参数，至于这些参数，就要根据实际开发时候的情况来进行更改。在当前情况下只需要更改上述两个参数即可，之后运行export.py，即可在/checkpoints目录下看到生成的onnx文件。

（图6 成功生成了yolov5s.onnx文件）

成功将Pytorch模型转为ONNX模型之后，我们需要验证这个方法是否是有效的，因此我们打开源码目录下的detect.py（图7），修改第244行中模型文件的路径，将其换成我们刚刚存放两个模型的路径

（图7 detect.py参数修改）

首先是对Pytorch模型进行测试，可见单张图片的推理速度在30ms左右。

（图8 Pytorch模型推理）

接着测试ONNX模型的推理速度，可见对于单张图片，单张图片推理速度提升至16ms左右，但是在pre-process环节又损失了很多速度。

（图9 onnx模型推理）

3. 基于C++的TensorRT开发

在本段中，我们将对基于C++的TensorRT开发进行简单描述。

首先，我们需要将我们刚才转出的ONNX模型导出为可以被TensorRT所调用的engine文件。我们打开我们下载好的TensorRT文件包，找到bin目录，我们可以见到目录下有一个trtexec.exe文件，接着我们将我们上一段中导出onnx文件复制到该目录下。

之后我们打开我们的终端或者命令提示符或者powershell，输入“trtexec.exe --onnx=yolov5s.onnx --saveEngine=yolov5s.engine”通常这一步需要一点时间，只要不提示Failed就耐心等待就好啦。

如上两图，我们已经成功创建了yolov5s.engine文件，我们就可以进行TensorRT模型推理了。

创建好engine文件之后，我们在VisualStudio（以下简称vs）中创建一个空白项目

省略一大堆写代码部分，我们直接贴出结果部分，具体代码实现可以参照GitHub仓库：

（图 TensorRT推理结果）

推理时间及资源消耗对比：

以上结果均为单张图片推理时间，在TensorRT推理中同样不包含数据的前处理以及输出的后处理如NMS等。

主函数部分源码解析：

在源码中，我们可以看到TensorRT推理的基本步骤

（1）在168-170行中，首先创建runtime、engine、context指针，这些指针分别对应运行时、engine模型、上下文初始化

（2）171行中对engine模型反序列化，意义是为了提前在GPU内存中获得地址

（3）172行中创建cuda流，目的是使用流来进行数据的传入、推理、结果传出

（4）后续就是定义显存消耗空间大小、模型推理等操作

在推理结束后，需要对使用的空间，主函数中实现如下：

首先是销毁流，其次将开辟的空间释放，接着删除内存上的指针，销毁所有GPU指针

到这里，关于TensorRT的C++推理也就结束了。

四、结语

在本文中，阐述了在算法开发工作中常见的几种给深度学习模型提升推理速度、减少资源消耗的方法，并且做了简单演示。虽然只是简单演示，但正是这些工作，能让AI更加贴近我们的生活。当然，针对硬件进行模型加速俨然成为了一种思想，在实际应用中，这种思想在现阶段深度学习模型的发展中已经完全进入到我们的生活之中，比如轻量化P图模型等，只是我们未曾感知罢了。期望在未来能有更多的人参与到深度学习算法开发的工作中来！

参考：[1]维基百科编者. ONNX[G/OL]. 维基百科, 2023(20230922)[2023-09-22]. .