NVIDIA TensorRT 其主要功能是通过处理网络定义并将其转换为优化引擎执行计划，加快深度学习推理。 TensorRT Engine Explorer (TREx) 是一个 Python 库和一组 Jupyter notebook，用于探索 TensorRT 引擎计划及其相关推理分析数据。

TREx 通过总结统计数据、图表实用程序和引擎图形可视化，为您提供生成引擎的可见性。 TREx 优化和调试高级网络性能非常有用，比如比较两个版本的网络性能。 深入的性能分析，NVIDIA Nsight Systems 推荐性能分析工具。

在这篇文章中，我总结了一下 TREx 工作数据和重点介绍工作流程 TensorRT 引擎的 API 功能。

TREx 主要抽象是 trex.EnginePlan，它包装了所有与发动机相关的信息。 EnginePlan 由多个输入 JSON 文件构成，每一个 JSON 文件描述引擎的不同方面，如数据依赖关系图及其分析数据。 EnginePlan 中间的信息可以通过 Pandas DataFrame 这是一种熟悉、强大、方便的数据结构。

在使用 TREx 之前，您必须构建和配置您的引擎。 TREx 脚本提供了一个简单的实用程序 process_engine.py 执行此操作。 作为参考，您可以选择任何方式收集此信息。

此脚本使用 trtexec 从 ONNX 模型构建引擎并分析引擎。 它还创造了几个 JSON 捕获引擎构建和分析会话的各个方面：

Plan-graph JSON 文件以 JSON 引擎数据流图的格式描述。

TensorRT 引擎Plan是 TensorRT 引擎的序列化格式。它包含有关最终推理图的信息，并且可以反序列化以进行推理运行时执行。

TensorRT 8.2 引入了 IEngineInspector API，它提供了检查引擎层、配置和数据依赖的能力。 IEngineInspector 使用简单的 JSON 提供此信息的格式化架构。 JSON 文件是 TREx trex.EnginePlan 对象的主要输入是必要的。

分析 JSON 文件为每个引擎层提供分析信息。

trtexec 实现命令行应用程序 IProfiler 并生成一个接口 JSON 该文件包含各层的分析记录。如果您只是想调查引擎的结构而没有相关的分析信息，则该文件是可选的。

JSON 文件包含每个分析迭代的时间记录。

分析引擎，trtexec 多个执行引擎可以平稳地测量噪声。每个引擎执行的时间信息可以作为时间记录 JSON 在文件中，平均测量值报告为发动机延迟。本文件是可选的，在评估和分析会话质量时通常非常有用。

元数据 JSON 文件描述了发动机的构造器配置和发动机的构造 GPU 的信息。 这一信息为引擎分析会话提供了更有意义的上下文，在比较两个或两个以上的引擎时特别有用。

TREx 工作流程 下图 总结了 TREx 工作流程：

• 首先，将你的深度学习模型转化为 TensorRT 网络。
• 构建和分析引擎，同时生成附带引擎 JSON 文件。
• 启动 TREx 探索文件的内容。

收集所有分析数据后，您可以创建一个 EnginePlan 实例：

plan = EnginePlan(     "my-engine.graph.json",     "my-engine.profile.json",     "my-engine.profile.metadata.json") 

使用 trex.EnginePlan 例子，你可以通过 Pandas DataFrame 访问大部分信息。 DataFrame 每一行代表计划文件中的一层，包括其名称、策略、输入、输出和其他属性的描述。

# Print layer names plan = EnginePlan("my-engine.graph.json") df = plan.df print(df['Name']) 

使用 DataFrame 抽象引擎信息非常方便，因为它有很多 Python 开发人员知道和喜欢 API，又强大了 API，具有切片、切块、导出、图形和打印数据的功能。

例如，列出引擎中最慢的三层非常简单：

# Print the 3 slowest layers top3 = plan.df.nlargest(3, 'latency.pct_time') for i in range(len(top3)):     layer = top3.iloc[i]    print("%s: %s" % (layer["Name"], layer["type"]))
features.16.conv.2.weight + QuantizeLinear_771 + Conv_775 + Add_777: Convolution
features.15.conv.2.weight + QuantizeLinear_722 + Conv_726 + Add_728: Convolution
features.12.conv.2.weight + QuantizeLinear_576 + Conv_580 + Add_582: Convolution

我们经常想对信息进行分组。 例如，您可能想知道每种层类型消耗的总延迟：

# Print the latency of each layer type
plan.df.groupby(["type"]).sum()[["latency.avg_time"]]

Pandas 可以与其他库很好地混合，例如 dtale，一个用于查看和分析数据框的便捷库，以及 Plotly，一个具有交互式绘图的图形库。这两个库都与示例 TREx 笔记本集成，但还有许多用户友好的替代方案，例如 qgrid、matplotlib 和 Seaborn。

还有一些方便的 API，它们是 Pandas、Plotly 和 dtale 的瘦包装器：

• 绘制数据（plotting.py）
• 可视化引擎图 (graphing.py)
• 交互式notebook（interactive.py 和 notebook.py）
• 报告（report_card.py 和 compare_engines.py）

最后，linting API (lint.py) 使用静态分析来标记性能危害，类似于软件 linter。理想情况下，层 linter 提供专家性能反馈，您可以根据这些反馈来提高引擎的性能。例如，如果您使用次优卷积输入形状或量化层的次优放置。 linting 功能处于早期开发状态，NVIDIA 计划对其进行改进。

TREx 还附带了一些教程notebook和两个工作流notebook：一个用于分析单个引擎，另一个用于比较两个或多个引擎。

使用 TREx API，您可以编写新的方法来探索、提取和显示 TensorRT 引擎，您可以与社区共享这些引擎。

现在您已经了解了 TREx 的运作方式，下面是一个展示 TREx 运行的示例。

在此示例中，您将创建一个量化的 ResNet18 PyTorch 模型的优化 TensorRT 引擎，对其进行分析，最后使用 TREx 检查引擎计划。 然后，您根据您的学习调整模型，以提高其性能。 此示例的代码可在 TREx GitHub 存储库中找到。

首先将 PyTorch ResNet 模型导出为 ONNX 格式。 使用 NVIDIA PyTorch 量化工具包在模型中添加量化层，但您不执行校准和微调，因为您专注于性能，而不是准确性。

在实际用例中，您应该遵循完整的量化感知训练 (QAT) 方法。 QAT 工具包会自动将伪量化操作插入 Torch 模型。 这些操作导出为 QuantizeLinear 和 DequantizeLinear ONNX 操作符：

import torch
import torchvision.models as models
# For QAT
from pytorch_quantization import quant_modules
quant_modules.initialize()
from pytorch_quantization import nn as quant_nn
quant_nn.TensorQuantizer.use_fb_fake_quant = True

resnet = models.resnet18(pretrained=True).eval()
# Export to ONNX, with dynamic batch-size
with torch.no_grad():
    input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(
        resnet, input, "/tmp/resnet/resnet-qat.onnx",
    	  input_names=["input.1"],
    	  opset_version=13,
    	  dynamic_axes={"input.1": {0: "batch_size"}})=

接下来，使用 TREx 实用程序 process_engine.py 脚本执行以下操作：

• 从 ONNX 模型构建引擎。
• 创建一个引擎计划 JSON 文件。
• 分析引擎执行并将结果存储在分析 JSON 文件中。 您还将计时结果记录在计时 JSON 文件中。

python3 <path-to-trex>/utils/process_engine.py /tmp/resnet/resnet-qat.onnx /tmp/resnet/qat int8 fp16 shapes=input.1:32x3x224x224

脚本 process_engine.py 使用 trtexec 来完成繁重的工作。 您可以通过简单地列出不带 -- 前缀的参数，从 process_engine.py 命令行透明地将参数传递给 trtexec。

在示例中，参数 int8、fp16 和 shape=input.1:32x3x224x224 被转发到 trtexec，指示它针对 FP16 和 INT8 精度进行优化，并将输入批量大小设置为 32。第一个脚本参数是输入 ONNX 文件（/tmp/resnet/resnet-qat.onnx），第二个参数（/tmp/resnet/qat）指向包含生成的JSON文件的目录。

您现在已准备好检查优化的引擎计划，因此请转到 TREx 引擎报告卡notebook。 我不会在这篇文章中介绍整个notebook，只介绍对这个示例有用的几个单元格。

第一个单元设置引擎文件并从各种 JSON 文件创建一个 trex.EnginePlan 实例：

engine_name = "/tmp/resnet/qat/resnet-qat.onnx.engine"
plan = EnginePlan(      f"{ 
          engine_name}.graph.json",      
  f"{ 
          engine_name}.profile.json",
  f"{ 
          engine_name}.profile.metadata.json")

下一个单元创建引擎数据依赖图的可视化，这对于理解原始网络到引擎的转换最有用。 TensorRT 将引擎作为拓扑排序的层列表执行，而不是作为可并行化的图。

默认的渲染格式是 SVG，它是可搜索的，在不同的比例下保持清晰，并且支持悬停文本以提供额外的信息而不占用大量空间。

graph = to_dot(plan, layer_type_formatter)
svg_name = render_dot(graph, engine_name, 'svg')

该函数创建一个 SVG 文件并打印其名称。 即使对于小型网络，在notebook内渲染也很麻烦，您可以在单独的浏览器窗口中打开 SVG 文件进行渲染。

TREx 图形 API 是可配置的，允许进行各种着色和格式化，并且可用的格式化程序包含信息。 例如，使用默认格式化程序，图层会根据其操作进行着色，并按名称、类型和分析延迟进行标记。 张量被描述为连接层的边缘，并根据它们的精度进行着色，并用它们的形状和内存布局信息进行标记。

在生成的 ResNet QAT 引擎图（下图）中，您会看到一些 FP32 张量（红色）。 进一步调查，因为您希望使用 INT8 精度执行尽可能多的层。 使用 INT8 数据和计算精度可提高吞吐量并降低延迟和功耗。

性能单元提供性能数据的各种视图，特别是每层精度视图（下图）显示了使用 FP32 和 FP16 进行的多个层计算。

report_card_perf_overview(plan)