基于TensorRT车辆实时推理优化

Optimizing NVIDIA TensorRT Conversion for Real-time Inference on Autonomous Vehicles

自动驾驶系统使用各种神经网络模型，这些模型要求在GPU上进行极其精确和高效的计算。Zoox是一家全新开发robotaxis的初创公司，充分利用了NVIDIA硬盘的高性能、节能计算功能。最近，Zoox在旧金山发布了一个一小时的全自动驾驶，详细展示了他们的AI堆栈。

与TensorFlow相比，NVIDIA TensorRT提供了显著的加速（fp32为2-6x，Zoox网络为9-19x），支持使用CUDA流的异步和并发推理能力。Zoox视觉/激光雷达/雷达/预测算法严重依赖于深部神经网络，这些神经网络在我们的车辆上都运行在NVIDIA gpu上，并且大多使用TensorRT部署。

TensorRT是一个用于高性能深度学习推理的SDK，它为深度学习推理应用程序提供低延迟和高吞吐量。

可以使用各种转换管道将模型转换为TensorRT引擎。例如，使用Caffe训练的模型可以使用Caffe解析器轻松地转换为TensorRT运行时。

但是，TensorFlow模型需要使用ONNX（开放式神经网络交换）转换为TensorRT引擎。本文中介绍的工具是针对TensorFlow的，但是这些原则也可以应用到其他的训练框架中。在为所有这些深层神经网络部署和维护TensorRT引擎的过程中，我们发现了以下痛点： ONNX和TensorRT只支持有限的TensorFlow操作集。

内核大小和步幅的某些组合可能会对TensorRT造成副作用。

迁移到精度降低的推理或TensorRT升级可能导致性能回归。

在Zoox，我们开发了一组工具来促进TensorRT引擎的部署、验证和维护，如图2所示。在下面的部分中，我们将详细介绍这些模块。

TensorRT conversion checker

TensorRT转换检查器的目标是帮助您在训练网络之前识别可能的转换失败。checker是轻量级的，按设计是最小的（在本文后面的代码示例中突出显示）。在训练之前，它会在所构建的网络上触发一个TensorRT转换过程。我们只有在转换成功后才开始训练。

Figure 2. Zoox TensorRT conversion pipeline overview.

下面的代码示例显示了TensorRT转换检查器。要使用插件，用户只需要导入数据包，在网络构建期间注册输入/输出节点，然后在训练开始之前触发转换检查。

import trt_checker

class Lenet5():

def network(self, X):

input = tf.identity(X, name = "input")

# Registers the input in the conversion checker.

trt_checker.register_input(input)

# Network definition.

...

# Output node.

output = tf.identity(logits, name="output")

# Registers the output node in the conversion checker.

trt_checker.register_output(output)

return output

def main():

...

# Checks if the model can be converted to trt.

conversion_result = trt_checker.check_conversion()

# Only train when trt conversion is successful.

if conversion_result:

accuracy = lenet_network.train()

Output deviation inspection

此插件的目标是在运行整个特定于模型的计算之前，报告转换的TensorRT引擎的潜在精度回归。这个插件在转换后的TensorRT引擎和原始TensorFlow图上运行推理，输入完全相同（由用户随机生成或指定）。然后报告输出偏差的分布，给开发人员一个潜在精度回归的预警。此模块是逐层检查模块的构建块。

Figure 3. Output deviation inspection

Layer-by-layer inspection

下面的代码示例显示逐层检查：

def layer_by_layer_analysis(graph, input_layer):

median_error = []

for layer in graph.layers():

errors = convert(graph, input=input_layer, output=layer)

median_error.append(median(errors))

plot(median_error)

如果观察到精度回归，我们希望找出TensorRT引擎中的哪个层或操作对回归有显著贡献。这促使我们开发了逐层检测模块。调用时，模块为每个中间操作运行一个转换作业，并报告此特定操作生成的中值/最大错误（如图4所示）。当研究在不同版本的TensorRT中观察到的不同行为时，此模块非常有用。

Figure 4. Example regression observed in semantic segmentation when upgrading from TensorRT 5.1.5 to TensorRT 7.0.

图5显示了这种回归的一个例子，在这个例子中，我们观察到语义分割输出有轻微的回归。我们对TensorRT 5.1引擎和TensorRT 7.0引擎进行了逐层检查，然后绘制了每层的中间误差。

图5显示了每一层产生的中值误差。我们可以看到在这个特定网络的上采样层中可能存在一个bug。基于这些信息，我们能够在一个较小的网络上重现这种回归，并将此错误报告给NVIDIA。此错误现在已在TensorRT 7.1中修复。

图5. 图4中使用的两台TensorRT发动机的逐层检查结果。橙色线显示了TensorRT 7.0推理输出与TensorFlow推理输出相比的中间误差。蓝线显示了TensorRT 5.0引擎生成的结果。上采样层的误差分布有显著差异。

Latency flame graph

为了可视化推理中的瓶颈并找出可能的优化操作，我们将TensorRT剖面仪生成的逐层计时信息绘制成火焰图。计时详细信息根据每个层的名称范围分组，如图6所示。这使我们能够看到网络的哪个部分花费的时间比预期的要长。

Figure 6. Latency flame graph on Inception Net. The 1767 samples shown in this graph indicates that a forward pass on this network takes 1.767 ms.

Automated conversion pipeline

在Zoox，我们维护一个自动转换管道，跟踪每个模型使用的转换选项。当触发时，自动转换管道将所有记录的模型转换为TensorRT引擎，并将它们上载到云中进行部署。它还为新转换的TensorRT引擎运行验证作业以验证准确性。这个管道帮助我们只用一个命令就可以将所有现有的模型升级到TensorRT的新版本。

Incompatible graph test suite

Zoox维护一个TensorFlow到TensorRT的转换测试套件。它测试了从TensorFlow图到tensort引擎的转换失败案例，以及报告的NVIDIA错误识别。

每个测试构建一个TensorFlow图，将其转换为TensorRT，并将输出偏差与TensorFlow图进行比较。使用此测试套件，我们不仅可以向Zoox工程师演示哪些图形结构或操作可能不适用于TensorRT，而且还可以检测升级到TensorRT的新版本时修复了哪些回归。

Summary

在这篇文章中，我们介绍了Zoox-TensorRT转换管道中的几个特性。TensorRT转换检查器参与神经网络训练的早期阶段，以确保在您浪费时间和资源进行全面训练之前发现不兼容的操作。可以在每一层调用推理精度验证来识别不利于降低精度计算的操作。详细的分析揭示了不必要的计算，这些计算不是在TensorRT内部优化的，但是可以通过在图构建期间简单的代码更改来优化。

自动转换管道帮助您验证每个TensorRT升级或模型重新转换。利用该流水线，我们成功地为神经网络在Zoox自主驾驶平台上执行各种流线型感知任务提供了TensorRT转换支持。