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 研习社编译整理。
雷锋网版权文章,未经授权禁止转载。详情见。