日前,Google 和辉达宣布将 NVIDIA TensorRT 整合到 TensorFlow 1.7。Google 开发者部落格介绍了此次合作的详细资讯及整合后的效能。
ensorRT 是一个可以用于最佳化深度学习模型,以进行推理,并为生产环境的 GPU 建立执行环境的程式库。它能最佳化 TensorFlow 的 FP16 浮点数和 INT8 整型数,并能自动选择针对特定平台的核心,以最大化吞吐量,并最大限度的降低 GPU 推理期间的延迟。全新的整合工作流程简化了 TensorFlow 使用 TensorRT 的步骤,同时使 TensorFlow 达到世界一流的效能水准。
经测试,在 NVIDIA Volta Tensor 核心,整合了 TensorRT 的 TensorFlow 执行 ResNet-50 比没有整合 TensorRT 的 TensorFlow 执行速度提高了 8 倍。
最佳化 TensorFlow 的子图
ensorFlow 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 最佳化,并用一个单一节点代替了(图中 highlight 部分)。
经过最佳化的 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 模型推理,能提供相较于单精准度模型 8 倍的吞吐量。相较于更高精准度的 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
详细的安装说明可在这里找到。
- Announcing TensorRT integration with TensorFlow 1.7
(本文由 雷锋网 授权转载;首图来源:TensorFlow)
延伸阅读:
- Google 发表 TensorFlow 1.5,全面支援动态图机制和 TensorFlow Lite
- Nvidia 发表 TensorRT 3 可程式化推理加速器,比起 CPU 能实现高达 40 倍吞吐率