对使用TensorRT加速AI深度学习推断效率的探索

施一飞

【摘 要】本文对使用TensorRT对目前流行的人工智能框架训练的神经网络模型的优化做了初步探索。目前基于人工智能的应用呈现爆炸式增长，图像识别、语音识别，自然语言处理、图像检索等功能已经成了人们日常必须的工具，随之而来对人工智能的推理效率和响应速度则提出更高要求。本文结合英伟达（nVidia）的TensorRT 3引擎和面向无人智能化的Jetson TX2平台，对以上问题的解决做了初步尝试。

【关键词】人工智能；TensorRT；TensorFlow；神经网络

作为人工智能浪潮中数字产品和服务的消费者，我们每天都会与多种人工智能相关的服务进行交互，包括语音识别，语言翻译，图像识别或视频字幕生成等。在这些产品和服务的背后，神经网络计算每个功能请求的结果。这个步骤在人工智能领域通常被称为“推理”，即新的数据通过传递给受过训练的神经网络计算以產生结果。在过往的机器学习专业中，有时称这个过程为“预测”或“计分”。

目前流行的做法是将神经网络部署在云端Web服务中，接收来自数千或数百万用户的并发请求，计算每个请求的推理结果并响应给用户。为了提供良好的用户体验，在考虑包括网络延迟，神经网络执行消耗和基于生产环境的其他延迟的情形下，努力促成最少响应时间。业界主要通过GPU的多核心并发计算能力和浮点运算的优势，利用诸如Caffe、TensorFLow等深度学习框架加速深度神经网络的训练，解决特定应用程序的效率要求。这些框架可以进行网络设计，执行模型评估和诊断以及用新数据重新训练模型，为快速构建问题解决原型提供了很大的灵活性。

1 深度学习训练和部署

用深度神经网络解决机器学习的监督问题涉及两个步骤：第一步，是使用GPU对海量标记数据进行深度神经网络训练。这一步中，神经网络将学习数以百万计的权重或参数，使其能够在示例数据与正确的响应之间建立映射。因为目标函数相对于网络权重已经被最小化，训练过程通常需要通过网络实现前推回代。通常，模型训练和精确度需要根据非训练期间的数据进行验证，以预估在实景的表现。第二步，是使用训练好的模型根据新采集的数据来进行预测。该步骤中，训练效果最好的模型被用于运行在生产环境中的应用程序，如数据中心，汽车或嵌入式平台。对于某些应用，如自动（辅助）驾驶，推理预测是需要实时完成的，那么计算中对高吞吐量、响应时间的驾驭至关重要。

如果目标是使用训练后的神经网络来感知其周围的嵌入式设备，那么通过模型的前向推断直接影响整体响应时间和设备消耗的功率。其中，要优化的关键指标是功率效率，即每瓦推断性能。每瓦性能也是数据中心运营效率最大化的关键指标。在这种情况下，需要尽量减少地理上和时间上大量不相关的请求所引起的延迟和消耗的能量，这往往也成了限制大批量计算能力的关键因素。

2 用TensorRT优化训练后神经网络

TensorRT通过优化训练后的神经网络来生成出于待部署状态的推理引擎，能最大限度地提高GPU推理性能和功效，尤其可以为诸如图像分类、分割和对象检测之类的深度学习应用提供最大的推理吞吐量和效率。TensorRT运行时，对从P4和V100等数据中心GPU到自动驾驶和嵌入式平台（如Drive PX2和Jetson TX2）的每个GPU平台的依赖性最小。

使用TensorRT先后需要完成编编译和部署两个阶段任务。在编译阶段，TensorRT对网络配置进行优化，生成优化计划，用于计算通过深度神经网络的正向推理。部署时采用驻守服务或用户应用程序的形式接收输入的批量数据，推理时根据输入数据执行优化计划，最后返回输出数据（分类，对象检测等）。

这里以TensorRT部署TensorFlow模型为例，说明如何在TensorRT中导入和优化经过训练的TensorFlow神经网络并将其部署为TensorRT运行时引擎。需要完成两个步骤：（1）导入并优化训练过的模型（这里需要使用TensorFlow框架的freeze_graph工具导出为冻结模型再导入），设置部署目标使用的GPU并发量、进度和内存等优化选项，生成推理引擎。以文件的形式存储在磁盘上。（2）加载和反序列号把存好的推理引擎计划文件，创建TensorRT引擎对象，使用它来进行目标部署平台上的新数据推理。注意，这里我们可以不用安装深度学习框架就可完成。下面是示例代码：

代码中UFF代表通用框架格式（Universal Framework format），它是运行优化引擎前用于表示网络模型的TensorRT内部格式。函数from_tensorflow_frozen_model（）负责进行模型转换，第一个参数是待导入的TensorFlow训练模型。第二个参数是输出图层名称。目前在TensorRT中支持一般的标准图层类型，包括卷积层、缩放层、完全连接层等。上述代码的输出是TensorFlow模型的UFF图形表示，它可以通过TensorRT进行分析。我们通过提供输入层的名称和维度（CHW格式）和输出层的名称来配置UFF解析器。接下来的编译过程，需要在目标平台进行编译，这里我们使用的是nVidia Jetson TX2，在需要使用Jetson TX2完成编译。在设置TensorRT的优化参数时，设定推理并发量、内存占用和目标进度：

参数指定并发量为1，使用1MB的暂存空间进行进度为FP32的推断。函数的输出就是TensorRT优化的运行时引擎，可以用来进行推理。优化的内容包括图层和张量垂直融合，内核的自动调优，张量动态记忆等。最后通过trt.utils.write_engine_to_file函数将引擎序列化为文件存储。当需要执行TensorRT的推断时，加载序列化的模型，进行推理。

而经过测试，在Jetson TX2平台上，用GPU运行TensorRT引擎优化后的TensorFlow模型，推理速度提高了近4倍。可见，合理有效的利用TensorRT，可以极大的优化目前嵌入式或汽车GPU平台的AI推理速度和响应时间。

【参考文献】

[1]NVIDIA CUDA Programming Guide. http：//docs.nvidia.com/cuda/pdf/CUDA_C_Programming_Guide.pdf.2013.

[2]面向GPU计算平台的若干并行优化关键技术研究[D]. 贾海鹏.中国海洋大学，2012.endprint