本文翻译于博客Speed up TensorFlow Inference on GPUs with TensorRT,这篇博客介绍了如何使用TensorRT加速TensorFlow模型的推理速度,作者为:
TensorFlow是当今最受欢迎的深度学习框架,在全球有着数十万名用户。NVIDIA®TensorRT™是一个深度学习平台,用于优化神经网络模型,加快数据中心,嵌入式芯片和汽车设备中运行的GPU加速平台的推理速度。我们对TensorFlow与TensorRT的集成感到兴奋,这似乎很自然,特别是NVIDIA提供的平台非常适合用于加速TensorFlow。这使得TensorFlow用户在使用TensorRT时具有极高的推理性能和近乎透明的工作流程。
TensorRT对神经网络图执行几项重要的转换和优化(图2)。首先,消除具有未使用的输出层以避免不必要的计算。接下来,在可能的情况下,对卷积,偏置和ReLU层进行融合以形成单个层。另一种转换是水平层融合或层聚合,以及聚合层到它们各自输出的划分。水平层融合通过对使用相同源tensor的层进行组合,并应用具有相似参数的相同操作来提高性能。请注意,这些图优化操作不会更改计算图中的基础计算:相反,它们会对计算图进行重新构建,使其可以更快,更有效地进行推理。
如果您已经将TensorRT与TensorFlow模型一起使用,那么您应该知道,在过去的版本中,使用TensorRT优化需要导出训练好的TensorFlow计算图。您还需要手动导入某些不受支持的TensorFlow图层,然后在TensorRT中运行完整图形。在大多数情况下,您再也不用这样做了。在新的工作流中,您可以通过使用简单的API在TensorFlow中应用TensorRT强大的FP16和INT8优化。现有的TensorFlow程序只需要几行新代码就可以使用这些优化。
在ResNet-50的基准测试钟,TensorRT将TensorFlow推断速度提高了8倍。这些性能改进仅需几行额外代码,并可与TensorFlow 1.7及更高版本一起使用。在本文中,我们将介绍新的工作流程和API,以帮助您开始使用它。
如图3所示,在TensorFlow推理工作流中加入tensorRT优化需要一个额外操作,在这个额外操作(使用绿色进行高亮)中,TensorRT根据frozen TensorFlow graph构建优化后的推理图。
为了完成优化操作,TensorRT对frozen TensorFlow graph进行解析,选出图中可以进行优化的子图。之后,TensorRT对子图进行优化,并且在原TensorFlow graph中将需要优化的子图替换为TensorRT节点,图中其余部分保持不变。在推理过程中,TensorFlow将调用TensorRT来执行优化后的TensorRT节点。通过这种方法,开发人员可以继续使用灵活的TensorFlow功能集和TensorRT优化。
举个栗子,假设有个计算图由A,B,C三个子图组成,TensorRT对子图B进行了优化,之后将子图B替换为TensorRT节点。在推理过程中,TensorFlow执行了子图A,之后调用TensorFlow执行了子图B,最后使用TensorFlow执行了子图C。
TensorRT优化了TensorFlow图中可能的最大子图。子图中的计算越多,从TensorRT获得的收益越大。您希望对大部分计算图进行优化,并使用最少数量的TensorRT节点替换,以获得最佳性能。根据计算图中的操作,优化后的计算图可能包含多个TensorRT节点。使用TensorFlow API,您可以指定子图中节点的最小数量,以便将其转换为TensorRT节点。任何小于指定设定节点数的子图都不会转换为TensorRT引擎,即使其与TensorRT兼容。这对于包含由不兼容节点分隔的小兼容子图的模型非常有用,从而导致微小的TensorRT引擎。
让我们看看如何更详细地实现工作流。
新的TensorFlow API支持使用几行新代码直接实现TensorRT优化。首先,指定分配给TensorFlow的GPU内存比例,剩余内存将用于TensorRT引擎。这可以通过使用GPUOptions
函数的新参数per_process_gpu_memory_fraction
来完成。需要在TensorFlow-TensorRT进程第一次启动时设置此参数。例如,将per_process_gpu_memory_fraction
设置为0.67会为TensorFlow分配67%的GPU内存,为TensorRT引擎分配剩余三分之一的GPU内存。
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction =
0 < memory_for_TensorFlow < 1)
下一步是使用tensorRT对TensorFlow计算图进行分析,进行优化,并将需要优化的子图替换为tensotRT节点。您可以使用新函数create_inference_graph
来对frozen TensorFlow graph进行tensorRT优化,这个函数的输入为frozen TensorFlow graph,返回含有tensorRT节点的优化图,如下面的代码片段所示:
trt_graph = trt.create_inference_graph(
input_graph_def = frozen_graph_def,
outputs = output_node_name,
max_batch_size=batch_size,
max_workspace_size_bytes=workspace_size,
precision_mode=precision,
minimum_segment_size=3)
让我们看看这个函数的参数:
input_graph_def
:frozen TensorFlow graphoutputs
:输出节点的tensor name,例如:[“resnet_v1_50/predictions/Reshape_1
”]max_batch_size
:整数,输入batch的大小,例如:16max_workspace_size_bytes
:整数,tensorRT可用的最大内存precision_mode
:string,可选值为:“FP32”, “FP16” 或者 “INT8”minimum_segment_size
:整数(默认为3),控制要创建的TensorRT引擎的子图中的最小节点数,如果子图中的节点数小于这个值,那么这个子图将不会被替换为tensorRT节点per_process_gpu_memory_fraction
和per_process_gpu_memory_fraction
两个参数应该被一起使用,对TensorFlow和TensorRT的内存分配进行控制,进而获得应用程序的最佳性能。为了最大化推理性能,您可以分配给tensorRT多一点内存,将剩余的内存分配给TensorFlow。例如:在一块内存为12G的GPU中,如果您将per_process_gpu_memory_fraction
设置为( 12–4 ) / 12 = 0.67,那么参数max_workspace_size_bytes
应设置为4000000000,为TensorRT引擎收集大约4G的内存。同样,最佳的内存分割方案是依赖于应用程序的,可能需要一些迭代来寻找最佳划分。
在对 ResNet-50进行TensorRT优化之后,TensorBoard可以使我们清楚地看到ResNet-50计算图中的变化。如图4所示,TensorRT几乎优化了整张计算图,并使用名称为“my_trt_op0”(使用红色进行高亮)的单个节点进行替换。根据模型中的层和操作,由于优化操作,TensorRT节点替换掉了图中的一部分。命名为“ conv1”的操作其实不是真正的卷积层,只是将feature map的格式由NHWC转换为NCHW(注:NCHW分别对应batch_size, channel, height, width)。
与FP32或FP64相比,使用半精度(也被称为FP16)可以降低神经网络的内存使用。FP16支持部署更大的网络,同时比FP32或FP64花费更少的时间。NVIDIA的Volta架构硬件采用了称为Tensor Cores的矩阵数学加速器。Tensor Cores提供4x4x4矩阵处理阵列,用于执行操作D = A * B + C,其中A, B, C and D为4×4的矩阵。图5展示了工作过程。矩阵乘法输入A和B是FP16矩阵,而累加矩阵C和D可以是FP16或FP32矩阵。
当使用FP16 math检测到推理时,TensorRT会自动使用硬件Tensor Cores。 Tensor Cores在NVIDIA Tesla V100上的峰值性能比双精度(FP64)快一个数量级,而吞吐量比单精度(FP32)提高了4倍。可以在函数create_inference_graph
中,将precision_mode
设置为“ FP16”启用半精度计算,如下所示。getNetwork()
用于从protobuf文件中读出frozen network,返回数据结构为tf.GraphDef()
的神经网络。
trt_graph = trt.create_inference_graph(
getNetwork(network_file_name),
outputs,
max_batch_size=batch_size,
max_workspace_size_bytes=workspace_size,
precision_mode=”FP16")
图6所示,在相同的硬件环境,小于7ms延时的情况下,应用基于NVIDIA Volta Tensor Cores的TensorFlow-TensorRT集成后的ResNet-50比仅使用TensorFlow的推理速度快8倍。
使用INT8精度执行推理进一步提高了计算速度并降低了对带宽的要求。减小的动态范围给表示神经网络的权重和激活带来挑战。
TensorRT可以将以单精度(FP32)或者半精度(FP16)训练的模型转化为以INT8量化部署的模型,同时可以最小化准确率损失。使用TensorRT优化对模型进行INT8量化之前需要对精度为FP32的模型进行校验。在创建TensorRT优化推理图之前,工作流程将包含校准步骤(Create Calibration Graph and Calibrate),如图7所示:
首先使用create_inference_graph
函数,将参数precision_mode
设置为INT8来校准模型,此功能的输出是一个frozen TensorFlow graph,可以进行校准。
trt_graph = trt.create_inference_graph(
getNetwork(network_file_name),
outputs,
max_batch_size=batch_size,
max_workspace_size_bytes=workspace_size,
precision_mode=”INT8")
现在使用校准数据运行校准图。 TensorRT使用节点数据的分布来量化节点的权重。**您必须使用与实际问题数据集分布相同或者类似的校准数据。**我们建议在首次使用INT8校准的模型时,检查推理过程中的误差累积。minimum_segment_size
可以帮助调整优化来最小化量化误差。通过使用minimum_segment_size
,您可以更改优化的INT8引擎中的最小节点数,进而更改最终优化计算图以微调结果准确性。
在校验数据上执行完计算图之后,使用函数calib_graph_to_infer_graph
对检验之后的计算图应用TensorRT优化,这个函数同样会将TensorFlow中的子图替换为INT8优化后的TensorRT节点。函数的输出为可以用于推理的frozen TensorFlow graph。
trt_graph=trt.calib_graph_to_infer_graph(calibGraph)
在使用这两个命令之后,就可以以INT8的精度对您的模型进行推理了。
如果您想查看此处显示的示例,请查看运行这些示例所需的代码:
https://developer.download.nvidia.com/devblogs/tftrt_sample.tar.xz
我们希望通过将TensorRT与TensorFlow集成,进而在使用NVIDIA GPU时可以获得最高性能,同时保持TensorFlow的易用性和灵活性。 NVIDIA将继续与Google TensorFlow团队密切合作,进一步增强这些集成功能。随着TensorRT支持更多网络,开发人员将自动受益于更新,而无需更改现有代码。
查找有关今天如何开始的说明:
https://www.tensorflow.org/install/install_linux
在不久的将来,我们希望标准的pip安装过程也能正常运行。敬请关注!
我们相信,在使用GPU时,您会看到将TensorRT与TensorFlow集成的实质性好处。您可以在以下位置找到有关TensorFlow的更多信息:https://www.tensorflow.org/.
有关TensorRT的更多信息,请访问NVIDIA的TensorRT页面https://developer.nvidia.com/tensorrt.