摘要TensorRT与TensorFlow的深度集成为GPU推理带来显著性能提升。通过新增API用户仅需几行代码即可实现FP16/INT8优化在ResNet-50测试中推理速度提升8倍。TensorRT自动执行计算图优化包括层融合和剔除同时保留TensorFlow的灵活性。Volta架构的张量核心可加速FP16运算而INT8校准流程能有效降低量化误差。该方案已在TensorFlow 1.7版本中实现未来将通过标准pip安装进一步简化部署流程。目录借助 TensorRT 提升 TensorFlow 在 GPU 上的推理速度为 TensorFlow 计算图应用 TensorRT 优化运用全新的 TensorFlow API借助 TensorBoard 可视化优化后的计算图在 Volta 架构 GPU 上运用张量核心Tensor Cores基于 INT8 精度的推理可用渠道借助 TensorRT 提升 TensorFlow 在 GPU 上的推理速度概述TensorFlow 仍是目前最受欢迎的深度学习框架在全球拥有数万名用户。英伟达 TensorRT 是一款深度学习平台可对神经网络模型进行优化提升其在数据中心、嵌入式设备和车载设备等 GPU 加速平台上的推理速度。我们对 TensorFlow 与 TensorRT 的集成感到十分欣喜二者的结合堪称天作之合尤其是英伟达还提供了适配性极佳的平台能够为 TensorFlow 实现加速。这一集成让 TensorFlow 用户既能获得极致的推理性能又能在使用 TensorRT 时拥有近乎无缝的工作流程。图 1TensorRT 对训练完成的神经网络模型进行优化生成可直接部署的运行时推理引擎。TensorRT 会对神经网络计算图执行多项关键的转换与优化操作见图 2。首先剔除输出未被使用的层避免无效计算其次在条件允许的情况下将卷积层、偏置层和 ReLU 激活层融合为单个层此外还会执行水平层融合也叫层聚合操作并将聚合后的层按需拆分至对应的输出节点。水平层融合通过合并接收相同源张量、且以相似参数执行相同操作的层实现性能提升。需要注意的是这些计算图优化不会改变图中的底层计算逻辑而是通过重构计算图让运算执行得更快、更高效。图 2a一个包含多个卷积层和激活层的卷积神经网络示例。bTensorRT 通过垂直层融合、水平层融合和层剔除优化简化了 GoogLeNet Inception 模块的计算图减少了计算量和内存开销。如果此前你曾在 TensorFlow 模型中使用过 TensorRT想必知道过去应用 TensorRT 优化时需要先导出训练好的 TensorFlow 计算图还需手动导入部分不被支持的 TensorFlow 层之后才能在 TensorRT 中运行完整的计算图。而现在绝大多数场景下都无需再执行这些操作。在新的工作流程中你只需通过简单的 API就能在 TensorFlow 内部借助 TensorRT 实现高效的 FP16半精度和 INT88 位整型优化。现有 TensorFlow 程序只需新增几行代码即可应用这些优化策略。在 ResNet-50 基准测试的低延迟运行场景中TensorRT 将 TensorFlow 的推理速度提升了 8 倍。这些性能提升仅需新增少量代码即可实现且适用于 TensorFlow 1.7 及更高版本。本文将为你介绍相关的新工作流程和 API助力你快速上手。为 TensorFlow 计算图应用 TensorRT 优化在 TensorFlow 推理工作流程中加入 TensorRT仅需新增一个步骤见图 3绿色高亮部分。该步骤中TensorRT 会基于冻结的 TensorFlow 计算图构建经过优化的推理计算图。图 3在 TensorFlow 中使用 TensorRT 的工作流程图具体实现时TensorRT 会解析冻结的 TensorFlow 计算图筛选出可优化的子图对其执行优化操作后将这些子图替换为原计算图中的 TensorRT 节点其余部分则保持不变。推理阶段TensorFlow 会执行完整的计算图在遇到 TensorRT 优化节点时调用 TensorRT 进行处理。通过这种方式开发者既能继续使用 TensorFlow 丰富灵活的功能集又能享受 TensorRT 带来的优化效果。我们以一个包含 A、B、C 三个段的计算图为例TensorRT 会对 B 段进行优化随后将其替换为单个节点。推理时TensorFlow 先执行 A 段再调用 TensorRT 执行 B 段最后执行 C 段。对用户而言操作方式与原先在 TensorFlow 中并无二致。TensorRT 会尽可能对 TensorFlow 计算图中规模最大的子图进行优化子图的计算量越大从 TensorRT 中获得的性能收益就越高。为实现最佳性能理想状态是让计算图的绝大部分都被优化并替换为尽可能少的 TensorRT 节点。根据计算图中的操作类型最终的计算图中可能会出现多个 TensorRT 节点。通过 TensorFlow API你可以指定子图转换为 TensorRT 节点所需的最小节点数任何节点数低于该阈值的子图即便与 TensorRT 兼容也不会被转换为 TensorRT 引擎。这一功能在处理这类模型时十分实用模型中包含多个由不兼容节点分隔的小型兼容子图若直接转换会生成多个微型 TensorRT 引擎影响性能。接下来我们详细介绍该工作流程的具体实现方法。运用全新的 TensorFlow API全新的 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_options tf.GPUOptions(per_process_gpu_memory_fraction 0 为TensorFlow分配的内存占比 1)下一步让 TensorRT 分析 TensorFlow 计算图、执行优化操作并将子图替换为 TensorRT 节点。可通过新的create_inference_graph函数对冻结的计算图应用 TensorRT 优化。该函数以冻结的 TensorFlow 计算图为输入返回包含 TensorRT 节点的优化后计算图代码示例如下trt_graph trt.create_inference_graph( input_graph_def frozen_graph_def, outputs output_node_name, max_batch_sizebatch_size, max_workspace_size_bytesworkspace_size, precision_modeprecision, minimum_segment_size3)下面为你介绍该函数的各项参数input_graph_def冻结的 TensorFlow 计算图outputs输出节点名称的字符串列表例如 [resnet_v1_50/predictions/Reshape_1]max_batch_size整型输入批次的大小例如 16max_workspace_size_bytes整型为 TensorRT 分配的最大 GPU 内存空间precision_mode字符串可选值为 “FP32”单精度、“FP16”半精度或 “INT8”8 位整型minimum_segment_size整型默认值为 3控制生成 TensorRT 引擎所需的子图最小节点数per_process_gpu_memory_fraction和max_workspace_size_bytes参数需配合使用合理划分 TensorFlow 与 TensorRT 的 GPU 内存才能实现应用的整体最佳性能。为最大化推理性能你可为 TensorRT 分配略多于其实际需求的内存剩余部分留给 TensorFlow。例如对于 12GB 的 GPU若将per_process_gpu_memory_fraction设为 (12-4)/120.67那么将max_workspace_size_bytes设为 4000000000即为 TensorRT 引擎分配了约 4GB 的内存。需要注意的是内存的最优分配比例与具体应用相关可能需要多次调试才能确定。借助 TensorBoard 可视化优化后的计算图通过 TensorBoard我们可以直观看到 ResNet-50 节点计算图在应用 TensorRT 优化后的变化。图 4 显示TensorRT 对几乎整个计算图都进行了优化并将其替换为一个名为 “my_trt_op0” 的节点红色高亮。根据模型中的层和操作类型TensorRT 会通过优化将模型的不同部分替换为 TensorRT 节点。图中的 “conv1” 并非实际的卷积层而是将数据格式从 NHWC 转换为 NCHW 的转置操作。图 4aTensorBoard 中的原始 ResNet-50 计算图 b应用 TensorRT 优化后子图被替换为 TensorRT 节点的 ResNet-50 计算图在 Volta 架构 GPU 上运用张量核心Tensor Cores与 FP32单精度或 FP64双精度相比使用 FP16半精度运算能减少神经网络的内存占用同时支持部署更大规模的网络且运算耗时更短。英伟达的 Volta 架构集成了名为张量核心的硬件矩阵数学加速器该核心包含一个 4×4×4 的矩阵处理阵列可执行运算D A * B C其中 A、B、C、D 均为 4×4 矩阵其工作原理见图 5。矩阵乘法的输入矩阵 A 和 B 为 FP16 格式而累加矩阵 C 和 D 可为 FP16 或 FP32 格式。图 5张量核心的矩阵处理运算原理当检测到使用 FP16 运算进行推理时TensorRT 会自动调用硬件张量核心。在英伟达 Tesla V100 显卡上张量核心的峰值性能比 FP64双精度高出约一个数量级吞吐量比 FP32单精度最高提升 4 倍。只需在create_inference_graph函数中将precision_mode参数设为 “FP16”即可启用半精度运算代码示例如下。其中getNetwork()为辅助函数用于从 protobuf 文件中读取冻结的网络并返回该网络的tf.GraphDef()对象。trt_graph trt.create_inference_graph( getNetwork(network_file_name), outputs, max_batch_sizebatch_size, max_workspace_size_bytesworkspace_size, precision_modeFP16)图 6 显示在相同硬件环境下结合 TensorFlow 与 TensorRT 并利用英伟达 Volta 张量核心运行 ResNet-50推理延迟可控制在 7 毫秒内速度相比单独运行 TensorFlow 提升了 8 倍。图 6ResNet-50 推理吞吐量性能对比基于 INT8 精度的推理采用 INT8 精度进行推理能进一步提升计算速度同时降低对带宽的要求。但 INT8 的动态范围更小难以精准表示神经网络的权重和激活值。针对这一问题TensorRT 提供了相应能力可将以 FP32单精度和 FP16半精度训练的模型转换为 INT8 量化格式部署同时将精度损失降至最低。将模型转换为 INT8 格式部署时需要在应用前文所述的 TensorRT 优化前对训练好的 FP32 模型进行校准工作流程也相应新增了校准步骤见图 7。图 7集成 INT8 推理的工作流程图首先调用create_inference_graph函数并将precision_mode参数设为 “INT8”对模型进行校准该函数的输出为可直接用于校准的冻结 TensorFlow 计算图trt_graph trt.create_inference_graph( getNetwork(network_file_name), outputs, max_batch_sizebatch_size, max_workspace_size_bytesworkspace_size, precision_modeINT8)随后使用校准数据运行该校准计算图TensorRT 会根据节点数据的分布特征对节点的权重进行量化。至关重要的是校准数据的分布需与实际生产环境中业务数据集的分布高度贴合。首次使用经 INT8 校准的模型时建议检查推理过程中的误差累积情况。minimum_segment_size参数可帮助调优优化后的计算图最大限度减少量化误差通过调整该参数改变 INT8 引擎优化子图的最小节点数进而调整最终的优化计算图实现推理结果精度的微调。在校准数据上执行完计算图后调用calib_graph_to_infer_graph函数对校准后的计算图应用 TensorRT 优化。该函数会将 TensorFlow 子图替换为针对 INT8 优化的 TensorRT 节点输出的冻结 TensorFlow 计算图可直接用于常规推理python运行trt_graphtrt.calib_graph_to_infer_graph(calibGraph)只需以上两条命令即可为你的 TensorFlow 模型启用 INT8 精度推理。若你想尝试本文中的示例可从以下地址获取相关运行代码https://developer.download.nvidia.com/devblogs/tftrt_sample.tar.xz可用渠道我们认为将 TensorRT 与 TensorFlow 集成后在使用英伟达 GPU 时能实现最优的推理性能同时保留 TensorFlow 的便捷性和灵活性。英伟达正与谷歌 TensorFlow 团队紧密合作进一步增强二者的集成能力。后续 TensorRT 将支持更多网络模型开发者无需修改现有代码就能自动享受相应的更新升级带来的收益。即刻上手的安装指南可参考https://www.tensorflow.org/install/install_linux我们预计在不久的将来通过标准的 pip 安装方式也可完成相关配置敬请期待我们相信在 GPU 环境下将 TensorRT 与 TensorFlow 集成你将获得显著的性能收益。更多关于 TensorFlow 的信息可访问其官方网站https://www.tensorflow.org/。更多关于 TensorRT 的信息可访问英伟达 TensorRT 专属页面https://developer.nvidia.com/tensorrt。