使用TREX探索NVIDIA TensorRT引擎

NVIDIA TensorRT 的主要功能是加速深度學(xué)習(xí)推理，通過(guò)處理網(wǎng)絡(luò)定義并將其轉(zhuǎn)換為優(yōu)化的引擎執(zhí)行計(jì)劃來(lái)實(shí)現(xiàn)。 TensorRT 發(fā)動(dòng)機(jī)瀏覽器 （ TREx ）是一個(gè) Python 庫(kù)和一組 Jupyter 筆記本，用于探索 TensorRT 引擎計(jì)劃及其相關(guān)推理評(píng)測(cè)數(shù)據(jù)。

TREx 提供了對(duì)生成引擎的可視性，通過(guò)匯總統(tǒng)計(jì)數(shù)據(jù)、圖表實(shí)用程序和引擎圖可視化為您提供了新的見(jiàn)解。 TREx 對(duì)于高級(jí)網(wǎng)絡(luò)性能優(yōu)化和調(diào)試非常有用，例如比較網(wǎng)絡(luò)兩個(gè)版本的性能。對(duì)于深入的性能分析，建議使用 NVIDIA Nsight Systems 進(jìn)行性能分析。

在這篇文章中，我總結(jié)了 TREx 的工作流程，并重點(diǎn)介紹了用于檢查數(shù)據(jù)和 TensorRT 引擎的 API 特性。要查看 TREx 的實(shí)際情況，我將通過(guò)《 insert action here 》完成如何實(shí)現(xiàn)《 value here 》的過(guò)程。

TREx 的工作原理

TREx 的主要抽象是trex.EnginePlan，它封裝了與引擎相關(guān)的所有信息。一個(gè)EnginePlan由幾個(gè)輸入 JSON 文件構(gòu)成，每個(gè)文件描述引擎的不同方面，例如其數(shù)據(jù)依賴(lài)關(guān)系圖和分析數(shù)據(jù)。EnginePlan中的信息可以通過(guò) Pandas 數(shù)據(jù)框訪(fǎng)問(wèn)，這是一種熟悉、強(qiáng)大且方便的數(shù)據(jù)結(jié)構(gòu)。

在使用 TREx 之前，必須構(gòu)建并分析引擎。 TREx 提供了一個(gè)簡(jiǎn)單的實(shí)用程序腳本 process_engine.py 來(lái)實(shí)現(xiàn)這一點(diǎn)。該腳本作為參考提供，您可以選擇任何方式收集此信息。

此腳本使用 trtexec 從 ONNX 模型構(gòu)建引擎并分析引擎。它還創(chuàng)建了幾個(gè) JSON 文件，用于捕獲引擎構(gòu)建和分析會(huì)話(huà)的各個(gè)方面：

平面圖 JSON 文件

計(jì)劃圖 JSON 文件以 JSON 格式描述引擎數(shù)據(jù)流圖。

TensorRT 引擎計(jì)劃是 TensorRT 引擎的序列化格式。它包含有關(guān)最終推理圖的信息，可以反序列化以執(zhí)行推理運(yùn)行時(shí)。

TensorRT 8.2 引入了 IEngineInspector API ，它提供了檢查引擎的層、層的配置及其數(shù)據(jù)依賴(lài)性的能力。IEngineInspector使用簡(jiǎn)單的 JSON 格式模式提供此信息。此 JSON 文件是 TREx trex.EnginePlan對(duì)象的主要輸入，是必需的。

分析 JSON 文件

分析 JSON 文件為每個(gè)引擎層提供分析信息。

trtexec 命令行應(yīng)用程序?qū)崿F(xiàn)了 IProfiler 接口和 生成 JSON 文件，其中包含每個(gè)層的分析記錄。如果您只想調(diào)查引擎的結(jié)構(gòu)，而不需要相關(guān)的分析信息，則此文件是可選的。

計(jì)時(shí)記錄 JSON 文件

JSON 文件包含每個(gè)分析迭代的計(jì)時(shí)記錄。

要對(duì)引擎進(jìn)行輪廓分析，trtexec多次執(zhí)行引擎以平滑測(cè)量噪聲。每個(gè)引擎執(zhí)行的計(jì)時(shí)信息可以作為單獨(dú)的記錄記錄在計(jì)時(shí) JSON 文件中，平均測(cè)量值報(bào)告為引擎延遲。此文件是可選的，通常在評(píng)估分析會(huì)話(huà)的質(zhì)量時(shí)非常有用。

如果您看到發(fā)動(dòng)機(jī)正時(shí)信息變化過(guò)大，可能需要確保您只使用 GPU ，并且計(jì)算和內(nèi)存時(shí)鐘已鎖定。

元數(shù)據(jù) JSON 文件

元數(shù)據(jù) JSON 文件描述了引擎的生成器配置以及用于構(gòu)建引擎的 GPU 的相關(guān)信息。此信息為引擎分析會(huì)話(huà)提供了更有意義的上下文，在比較兩個(gè)或多個(gè)引擎時(shí)尤其有用。

TREx 工作流

圖 1 總結(jié)了 TREx 工作流：

首先，將您的深度學(xué)習(xí)模型轉(zhuǎn)換為 TensorRT 網(wǎng)絡(luò)。

構(gòu)建和分析引擎，同時(shí)生成附帶的 JSON 文件。

旋轉(zhuǎn) TREx 以瀏覽文件的內(nèi)容。

NVIDIA TensorRT 的主要功能是加速深度學(xué)習(xí)推理，通過(guò)處理網(wǎng)絡(luò)定義并將其轉(zhuǎn)換為優(yōu)化的引擎執(zhí)行計(jì)劃來(lái)實(shí)現(xiàn)。 TensorRT 發(fā)動(dòng)機(jī)瀏覽器 （ TREx ）是一個(gè) Python 庫(kù)和一組 Jupyter 筆記本，用于探索 TensorRT 引擎計(jì)劃及其相關(guān)推理評(píng)測(cè)數(shù)據(jù)。

TREx 提供了對(duì)生成引擎的可視性，通過(guò)匯總統(tǒng)計(jì)數(shù)據(jù)、圖表實(shí)用程序和引擎圖可視化為您提供了新的見(jiàn)解。 TREx 對(duì)于高級(jí)網(wǎng)絡(luò)性能優(yōu)化和調(diào)試非常有用，例如比較網(wǎng)絡(luò)兩個(gè)版本的性能。對(duì)于深入的性能分析，建議使用 NVIDIA Nsight Systems 進(jìn)行性能分析。

在這篇文章中，我總結(jié)了 TREx 的工作流程，并重點(diǎn)介紹了用于檢查數(shù)據(jù)和 TensorRT 引擎的 API 特性。要查看 TREx 的實(shí)際情況，我將通過(guò)《 insert action here 》完成如何實(shí)現(xiàn)《 value here 》的過(guò)程。

TREx 的工作原理

TREx 的主要抽象是trex.EnginePlan，它封裝了與引擎相關(guān)的所有信息。一個(gè)EnginePlan由幾個(gè)輸入 JSON 文件構(gòu)成，每個(gè)文件描述引擎的不同方面，例如其數(shù)據(jù)依賴(lài)關(guān)系圖和分析數(shù)據(jù)。EnginePlan中的信息可以通過(guò) Pandas 數(shù)據(jù)框訪(fǎng)問(wèn)，這是一種熟悉、強(qiáng)大且方便的數(shù)據(jù)結(jié)構(gòu)。

在使用 TREx 之前，必須構(gòu)建并分析引擎。 TREx 提供了一個(gè)簡(jiǎn)單的實(shí)用程序腳本 process_engine.py 來(lái)實(shí)現(xiàn)這一點(diǎn)。該腳本作為參考提供，您可以選擇任何方式收集此信息。

此腳本使用 trtexec 從 ONNX 模型構(gòu)建引擎并分析引擎。它還創(chuàng)建了幾個(gè) JSON 文件，用于捕獲引擎構(gòu)建和分析會(huì)話(huà)的各個(gè)方面：

平面圖 JSON 文件

計(jì)劃圖 JSON 文件以 JSON 格式描述引擎數(shù)據(jù)流圖。

TensorRT 引擎計(jì)劃是 TensorRT 引擎的序列化格式。它包含有關(guān)最終推理圖的信息，可以反序列化以執(zhí)行推理運(yùn)行時(shí)。

TensorRT 8.2 引入了 IEngineInspector API ，它提供了檢查引擎的層、層的配置及其數(shù)據(jù)依賴(lài)性的能力。IEngineInspector使用簡(jiǎn)單的 JSON 格式模式提供此信息。此 JSON 文件是 TREx trex.EnginePlan對(duì)象的主要輸入，是必需的。

分析 JSON 文件

分析 JSON 文件為每個(gè)引擎層提供分析信息。

trtexec 命令行應(yīng)用程序?qū)崿F(xiàn)了 IProfiler 接口和 生成 JSON 文件，其中包含每個(gè)層的分析記錄。如果您只想調(diào)查引擎的結(jié)構(gòu)，而不需要相關(guān)的分析信息，則此文件是可選的。

計(jì)時(shí)記錄 JSON 文件

JSON 文件包含每個(gè)分析迭代的計(jì)時(shí)記錄。

要對(duì)引擎進(jìn)行輪廓分析，trtexec多次執(zhí)行引擎以平滑測(cè)量噪聲。每個(gè)引擎執(zhí)行的計(jì)時(shí)信息可以作為單獨(dú)的記錄記錄在計(jì)時(shí) JSON 文件中，平均測(cè)量值報(bào)告為引擎延遲。此文件是可選的，通常在評(píng)估分析會(huì)話(huà)的質(zhì)量時(shí)非常有用。

如果您看到發(fā)動(dòng)機(jī)正時(shí)信息變化過(guò)大，可能需要確保您只使用 GPU ，并且計(jì)算和內(nèi)存時(shí)鐘已鎖定。

元數(shù)據(jù) JSON 文件

元數(shù)據(jù) JSON 文件描述了引擎的生成器配置以及用于構(gòu)建引擎的 GPU 的相關(guān)信息。此信息為引擎分析會(huì)話(huà)提供了更有意義的上下文，在比較兩個(gè)或多個(gè)引擎時(shí)尤其有用。

TREx 工作流

圖 1 總結(jié)了 TREx 工作流：

首先，將您的深度學(xué)習(xí)模型轉(zhuǎn)換為 TensorRT 網(wǎng)絡(luò)。

構(gòu)建和分析引擎，同時(shí)生成附帶的 JSON 文件。

旋轉(zhuǎn) TREx 以瀏覽文件的內(nèi)容。

收集所有分析數(shù)據(jù)后，可以創(chuàng)建一個(gè)EnginePlan實(shí)例：

圖 1 ：。 TensorRT Engine Explorer 工作流

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

對(duì)于trex.EnginePlan實(shí)例，您可以通過(guò) pandasDataFrame對(duì)象訪(fǎng)問(wèn)大部分信息。數(shù)據(jù)框中的每一行表示計(jì)劃文件中的一個(gè)層，包括其名稱(chēng)、策略、輸入、輸出和描述該層的其他屬性。

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

使用數(shù)據(jù)幀抽象引擎信息很方便，因?yàn)樗仁窃S多 Python 開(kāi)發(fā)人員都知道并喜歡的 API ，也是一種功能強(qiáng)大的 API ，具有數(shù)據(jù)切片、切割、導(dǎo)出、繪圖和打印功能。

例如，列出引擎中三個(gè)最慢的層很簡(jiǎn)單：

# 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

我們經(jīng)常想把信息分組。例如，您可能想知道每種層類(lèi)型消耗的總延遲：

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

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"}})=

接下來(lái)，使用 TREx 實(shí)用程序process_engine.py腳本執(zhí)行以下操作：

從 ONNX 模型構(gòu)建引擎。

創(chuàng)建引擎計(jì)劃 JSON 文件。

分析引擎執(zhí)行并將結(jié)果存儲(chǔ)在分析 JSON 文件中。您還可以將計(jì)時(shí)結(jié)果記錄在一個(gè)計(jì)時(shí) JSON 文件中。

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

腳本process_engine.py使用trtexec來(lái)完成繁重的工作。您可以從process_engine.py命令行透明地將參數(shù)傳遞給trtexec，只需列出它們，而不需要--前綴。

在該示例中，參數(shù)int8、fp16和shapes=input.1:32x3x224x224被轉(zhuǎn)發(fā)到trtexec，指示其優(yōu)化 FP16 和 INT8 精度，并將輸入批次大小設(shè)置為 32 。第一個(gè)腳本參數(shù)是輸入 ONNX 文件（/tmp/resnet/resnet-qat.onnx），第二個(gè)參數(shù)（/tmp/resnet/qat）指向包含生成的 JSON 文件的目錄。

現(xiàn)在，您已經(jīng)準(zhǔn)備好檢查優(yōu)化的引擎計(jì)劃，所以請(qǐng)轉(zhuǎn)到 TREx 引擎報(bào)告卡筆記本 。在這篇文章中，我不會(huì)瀏覽整個(gè)筆記本，只有幾個(gè)單元格對(duì)這個(gè)例子有用。

第一個(gè)單元格設(shè)置引擎文件并創(chuàng)建 trex 。來(lái)自各種 JSON 文件的 EnginePlan 實(shí)例：

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")

下一個(gè)單元格創(chuàng)建引擎數(shù)據(jù)依賴(lài)關(guān)系圖的可視化，這對(duì)于理解原始網(wǎng)絡(luò)到引擎的轉(zhuǎn)換非常有用。 TensorRT 將引擎作為拓?fù)渑判虻膶恿斜韴?zhí)行，而不是作為可并行化的圖形執(zhí)行。

默認(rèn)呈現(xiàn)格式為SVG，可搜索，在不同比例下保持清晰，并支持懸停文本以提供附加信息，而不占用大量空間。

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

該函數(shù)創(chuàng)建一個(gè) SVG 文件并打印其名稱(chēng)。即使對(duì)于小型網(wǎng)絡(luò)，筆記本內(nèi)部的渲染也很麻煩，您可以在單獨(dú)的瀏覽器窗口中打開(kāi) SVG 文件進(jìn)行渲染。

TREx graphing API 是可配置的，允許使用各種顏色和格式，并且可用的格式設(shè)置程序包含信息。例如，使用默認(rèn)的格式化程序，層根據(jù)其操作進(jìn)行著色，并按名稱(chēng)、類(lèi)型和分析的延遲進(jìn)行標(biāo)記。張量被描述為連接各層的邊，并根據(jù)其精度進(jìn)行著色，并用其形狀和內(nèi)存布局信息進(jìn)行標(biāo)記。

在生成的 ResNet QAT 引擎圖（圖 3 ）中，您可以看到一些 FP32 張量（紅色）。進(jìn)一步研究，因?yàn)槟Ｍ褂?INT8 precision 執(zhí)行盡可能多的層。使用 INT8 數(shù)據(jù)和計(jì)算精度可以提高吞吐量，降低延遲和功耗。

圖 3 ：。 QAT ResNet18 引擎的數(shù)據(jù)依賴(lài)關(guān)系圖

性能單元提供了各種性能數(shù)據(jù)視圖，特別是每層精度視圖（圖 4 ）顯示了使用 FP32 和 FP16 計(jì)算的幾個(gè)層。

report_card_perf_overview(plan)

def forward(self, x: Tensor) -> Tensor: identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) if self._quantize: out += self.residual_quantizer(identity) else: out += identity out = self.relu(out) return out

更改 PyTorch 代碼后，必須重新生成模型，并使用修改后的模型再次遍歷筆記本單元格?，F(xiàn)在，您可以減少到三個(gè)重新格式化層，它們消耗了大約 20.5% 的總延遲（從 26.5% 下降），并且大多數(shù)層現(xiàn)在都以 INT8 精度執(zhí)行。

圖 7 ： QAT ResNet18 模式，在剩余連接上添加 Q / DQ 后

其余的 FP32 層圍繞網(wǎng)絡(luò)末端的全局平均池（ GAP ）層。再次修改模型以量化間隙層。

def _forward_impl（self， x： Tensor） -》 Tensor： # See note ［TorchScript super（）］ x = self.conv1（x） x = self.bn1（x） x = self.relu（x） x = self.maxpool（x） x = self.layer1（x） x = self.layer2（x） x = self.layer3（x） x = self.layer4（x） if self._quantize_gap： x = self.gap_quantizer（x） x = self.avgpool（x） x = torch.flatten（x， 1） x = self.fc（x） return x

使用新模型在筆記本單元中進(jìn)行最后一次迭代。現(xiàn)在只有一個(gè)重新格式化層，所有其他層都在 INT8 中執(zhí)行。搞定了！

圖 8 ：在剩余連接上添加 Q / DQ 并量化間隙層后的逐層精度視圖

現(xiàn)在您已經(jīng)完成了優(yōu)化，可以使用 發(fā)動(dòng)機(jī)比較筆記本 來(lái)比較這兩個(gè)引擎。此筆記本不僅在您正在積極優(yōu)化網(wǎng)絡(luò)性能時(shí)有用，而且在以下情況下也有用：

當(dāng)您想要比較為不同的 GPU HW 平臺(tái)或不同的 TensorRT 版本構(gòu)建的引擎時(shí)。

當(dāng)您想要評(píng)估層的性能如何跨不同的批處理大小進(jìn)行擴(kuò)展時(shí)。

了解發(fā)動(dòng)機(jī)之間的精度不一致是否是由于 TensorRT 層精度選擇不同所致。

發(fā)動(dòng)機(jī)比較筆記本提供了表格和圖形視圖來(lái)比較發(fā)動(dòng)機(jī)，這兩種視圖都適用，具體取決于您需要的詳細(xì)程度。圖 8 顯示了我們?yōu)?PyTorch ResNet18 模型構(gòu)建的五個(gè)引擎的疊加延遲。為簡(jiǎn)潔起見(jiàn)，我沒(méi)有討論創(chuàng)建 FP32 和 FP16 引擎，但這些引擎可以在 TREx GitHub 存儲(chǔ)庫(kù)中找到。

圖 9 ：同一 ResNet18 網(wǎng)絡(luò)的五個(gè)引擎的疊加延遲

為 FP16 精度優(yōu)化的引擎大約比 FP32 引擎快 2 倍，但也比我們首次嘗試的 INT8 QAT 引擎快。如前所述，這是由于許多 INT8 卷積輸出 FP16 數(shù)據(jù)，然后需要重新格式化層以顯式量化回 INT8 。

如果您只關(guān)注本文中優(yōu)化的三個(gè) QAT 引擎，那么您可以看到，在向剩余連接添加 Q / DQ 時(shí)，您是如何消除 11 個(gè) FP16 引擎層的。量化間隙層時(shí)，消除了另外兩個(gè) FP32 層。

圖 10 ：優(yōu)化的三臺(tái)發(fā)動(dòng)機(jī)的精度計(jì)數(shù)

您還可以查看優(yōu)化如何影響三個(gè)引擎的延遲（圖 10 ）。

圖 11 ：三個(gè)引擎的延遲，按層類(lèi)型分組

您可能會(huì)注意到一些看起來(lái)很奇怪的池層延遲結(jié)果：當(dāng)您量化剩余連接時(shí)，總池延遲下降了 10 倍，然后當(dāng)您量化間隙層時(shí)，總池延遲上升了 70% 。

這兩個(gè)結(jié)果都是違反直覺(jué)的，所以請(qǐng)仔細(xì)觀(guān)察它們。有兩個(gè)池層，第一次卷積之后是一個(gè)大的池層，最后一次卷積之前是一個(gè)小的池層。量化剩余連接后，可以使用 INT8 精度的輸出執(zhí)行第一個(gè)池和卷積層。它們與夾在中間的 ReLU 融合到 ConfactPool 層中，但浮點(diǎn)類(lèi)型不支持這種融合。

為什么間隙層在量化時(shí)延遲增加？這個(gè)層的激活大小很小，每個(gè) INT8 輸入系數(shù)都轉(zhuǎn)換為 FP32 ，以便使用高精度進(jìn)行平均。最后，將結(jié)果轉(zhuǎn)換回 INT8 。

該層的數(shù)據(jù)大小也很小，并且駐留在快速二級(jí)緩存中，因此額外的精度轉(zhuǎn)換計(jì)算相對(duì)昂貴。盡管如此，因?yàn)槟梢匀サ魢@間隙層的兩個(gè)重新格式化層，所以總的引擎延遲（這是您真正關(guān)心的）會(huì)減少。

總結(jié)

在這篇文章中，我介紹了 TensorRT 引擎瀏覽器，簡(jiǎn)要回顧了它的 API 和特性，并通過(guò)一個(gè)示例演示了 TREx 如何幫助優(yōu)化 TensorRT 引擎的性能。 TREx 可以在 TensorRT 的 GitHub 存儲(chǔ)庫(kù)中的 實(shí)驗(yàn)工具 目錄下找到。