Google 和 Nvidia 强强联手，带来优化版 TensorFlow 1.7

雷锋网 AI 研习社按，日前，谷歌和英伟达宣布将 NVIDIA TensorRT 集成到 TensorFlow 1.7 中。在谷歌开发者博客中，他们介绍了此次合作的详细信息以及整合之后的性能，雷锋网(公众号：雷锋网) AI 研习社编译整理如下：

TensorRT 是一个可以用于优化深度学习模型，以进行推理，并为生产环境中的 GPU 创建运行环境的库。它能优化 TensorFlow 中的 FP16 浮点数和 INT8 整型数，并能自动选择针对特定平台的内核，以最大化吞吐量，并最大限度的降低 GPU 推理期间的延迟。全新的集成工作流程简化了在 TensorFlow 中使用 TensorRT 的步骤，同时使得 TensorFlow 达到了世界一流的性能水平。

经测试，在 NVIDIA Volta Tensor 核心上，集成了 TensorRT 的 TensorFlow 运行 ResNet-50 比没有集成 TensorRT 的 TensorFlow 执行速度提高了 8 倍。

优化 TensorFlow 中的子图

在 TensorFlow 1.7 中，TensorRT 可以用于优化子图，而 TensorFlow 执行其余未优化的部分。这个方法使得开发者既能够使用 TensorFlow 的众多功能来快速构建模型，同时也可以在执行推理时使用 TensorRT 获得强大的优化能力。如果你尝试过在之前的 TensorFlow 模型中使用 TensorRT，你应该知道，要想使用某些不受支持的 TensorFlow 层，必须手动导入，这在某些情况下可能会耗费大量时间。

从工作流程的角度来看，开发者可以使用 TensorRT 来优化 TensorFlow 的每个子图。

在推断过程中，TensorFlow 先将执行所有支持区域的图，之后调用 TensorRT 去执行那些经过 TensorRT 优化过的节点。举个例子，如果你的图包含 A，B，C 三段，其中 B 段被 TensorRT 优化过，B 将被一个节点代替。那么在推理过程中，TensorFlow 将先执行 A，之后调用 TensorRT 执行 B，最后 TensorFlow 执行 C。

这个用于优化 TensorRT 的新加入的 TensorFlow API，以冻结的 TensorFlow 图为输入，针对该子图进行优化，最后将优化过的推理子图发送回 TensorFlow 中。

下面为一段示例代码：

# Reserve memory for TensorRT inference engine

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = number_between_0_and_1)

… 

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)  # Get optimized graph

per_process_gpu_memory_fraction 这个参数定义了 TensorFlow 允许使用的 GPU 显存的比例，剩余的显存将分配给 TensorRT。这个参数应该在 TensorFlow-TensorRT 进程第一次启动的时候设定好。比如，per_process_gpu_fraction=0.67，那么 67% 的显存会被分配给 TensorFlow，其余的 33% 会被分配给 TensorRT 引擎。

Create_inference_graph 函数将冻结住的 TensorFlow 图作为输入，返回一个经过 TensorRT 节点优化过的图。我们看看这个函数的参数：

• Input_graph_def:冻结住的 TensorFlow 图

• Outputs:输出节点名字的字符串列表，比如：[“resnet_v1_50/predictions/Resape_1”]

• Max_batch_size:整数，输入的 batch size，比如，16

• Max_workspace_size_bytes:整数，能分配给 TensorRT 的最大 GPU 显存大小

• Precision_mode:字符串，可选的值为「FP32」, 「FP16」, 「INT8」

举个例子，如果 GPU 有 12GB 显存，想要给 TensorRT 引擎分配 4GB 显存，那么应该设置 per_process_gpu_memory_fraction 为（12-4）/12=0.67，max_workspace_size_bytes=4,000,000,000.

我们来试着将这个新的 API 应用在 ResNet-50 上，看看经过优化后的模型在 TensorBoard 中看起来是什么样的。左侧的图像是没有经过 TensorRT 优化的 ResNet-50，右侧是经过优化的。在这个设定下，大部分图被 TensorRT 优化，并用一个单一节点代替了（图中高亮部分）。

经过优化的INT8推理性能

TensorRT 兼容单精度（FP32）和半精度（FP16）训练的模型（也可以将它们量化为 INT8），同时能尽可能减少由精度降低而导致的准确率降低。INT8 模型能够更快的计算，同时对带宽的需求也会降低，但是因为可用的动态范围降低了，这也对神经网络的权重和激活表示提出了很大的挑战。

为了解决这个问题，TensorRT 使用了一个校正过程，以尽可能减小将 FP32 网络近似成 8-bit 整型表示时的信息损失。在使用 TensorRT 优化 TensorFlow 图之后，可以使用下面的命令将图传递给 TensorRT 进行校准，如下：

trt_graph=trt.calib_graph_to_infer_graph(calibGraph)

除此之外的网络推理流程都没有变化。这一步的输出为一个可以被 TensorFlow 执行的冻结图。

在NVIDIA Volta GPU上自动使用Tensor核心

在 NVIDIA Volta GPU 的 Tensor 核心上通过 TensorRT 进行半精度 TensorFlow 模型推理，能够提供相较于单精度模型八倍的吞吐量。相较于更高精度的 FP32 或者 FP64，半精度数据（FP16）能够减少神经网络的显存使用量，这使得开发者能够训练和部署更大规模的神经网络，同时 FP16 相比 FP32 和 FP64 有更少的传输时间。

如果每个 Tensor 核心执行的是 D=A*B+C，其中 A 和 B 为半精度 4*4 矩阵，D 和 C 是单精度或者半精度 4*4 矩阵，那么 V100 上此时 Tensor 核心的峰值性能是双精度（FP64）性能的 10 倍，是单精度（FP32）性能的 4 倍。

Google 目前已经发布了 TensorFlow 1.7，同时也将跟 NVIDIA 更紧密地合作。希望这个新的解决方案额能够在提供最强性能的同时，保持 TensorFlow 的易用性和灵活性。随着 TensorRT 支持越来越多的网络结构，大家只要更新就可以享受到这些好处，而无须改写代码。

使用标准pip install即可更新到 TensorFlow 1.7:

pip install tensorflow-gpu r1.7

详细的安装说明可在下面找到：

https://github.com/tensorflow/tensorflow/tree/r1.7/tensorflow/contrib/tensorrt

via: Google Develop Blog

雷锋网 AI 研习社编译整理。

雷锋网版权文章，未经授权禁止转载。详情见。