使用TREX探索NVIDIA TensorRT引擎

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

TREx 提供了對(duì)生成引擎的可視性，通過匯總統(tǒng)計(jì)數(shù)據(jù)、圖表實(shí)用程序和引擎圖可視化為您提供了新的見解。 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í)際情況，我將通過《 insert action here 》完成如何實(shí)現(xiàn)《 value here 》的過程。

TREx 的工作原理

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

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

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

平面圖 JSON 文件

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

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

TensorRT 8.2 引入了 IEngineInspector API ，它提供了檢查引擎的層、層的配置及其數(shù)據(jù)依賴性的能力。IEngineInspector使用簡單的 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ì)話的質(zhì)量時(shí)非常有用。

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

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

元數(shù)據(jù) JSON 文件描述了引擎的生成器配置以及用于構(gòu)建引擎的 GPU 的相關(guān)信息。此信息為引擎分析會(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)容。

圖 1 TensorRT Engine Explorer 工作流

TREx 功能和 API

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

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

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

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

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

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

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

# 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)常想把信息分組。例如，您可能想知道每種層類型消耗的總延遲：

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

pandas 與其他庫很好地結(jié)合，如用于查看和分析數(shù)據(jù)幀的方便庫 dtale 和帶有交互式繪圖的圖形庫 Plotly 。這兩個(gè)庫都與示例 TREx 筆記本集成，但有許多用戶友好的 備擇方案 ，如 qgrid 、 matplotlib 和 Seaborn 。

還有一些方便的 API ，它們是 Pandas 、 Plotly 和 dtale 的薄包裝：

打印數(shù)據(jù)（plotting.py）

可視化引擎圖（graphing.py）

交互式筆記本（interactive.py和notebook.py）

報(bào)告（report_card.py和compare_engines.py）

最后， linting API （lint.py）使用靜態(tài)分析來標(biāo)記性能危害，類似于軟件 linter 。理想情況下，層過梁提供專家性能反饋，您可以根據(jù)這些反饋采取行動(dòng)，以提高發(fā)動(dòng)機(jī)的性能。例如，如果使用次優(yōu)卷積輸入形狀或次優(yōu)量化層放置。 linting 功能處于早期開發(fā)狀態(tài)， NVIDIA 計(jì)劃對(duì)其進(jìn)行改進(jìn)。

TREx 還附帶了兩個(gè)教程筆記本和兩個(gè)工作流筆記本：一個(gè)用于分析單個(gè)引擎，另一個(gè)用于比較兩個(gè)或多個(gè)引擎。

使用 TREx API ，您可以編寫新的方法來探索、提取和顯示 TensorRT 引擎，您可以與社區(qū)共享。

TREx 演練示例

現(xiàn)在您已經(jīng)了解了 TREx 的操作方式，下面是一個(gè)顯示 TREx 實(shí)際操作的示例。

在本例中，您創(chuàng)建了一個(gè)量化的優(yōu)化 TensorRT 引擎 ResNet18 PyTorch ，對(duì)其進(jìn)行分析，最后使用 TREx 檢查發(fā)動(dòng)機(jī)計(jì)劃。］然后根據(jù)所學(xué)內(nèi)容調(diào)整模型，以提高其性能。此示例的代碼可在 TREx GitHub 存儲(chǔ)庫中找到。

首先，將 PyTorch ResNet 模型導(dǎo)出為 ONNX 格式。使用 NVIDIA PyTorch 量化工具包 用于在模型中添加量化層，但您不執(zhí)行校準(zhǔn)和微調(diào)，因?yàn)槟P(guān)注的是性能，而不是準(zhǔn)確性。

在實(shí)際用例中，您應(yīng)該遵循完整的量化感知訓(xùn)練（ QAT ）方法。 QAT 工具包自動(dòng)將假量化操作插入火炬模型。這些操作導(dǎo)出為 QuantizeLinear 和 DequantizeLinear ONNX 運(yùn)算符：

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 實(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來完成繁重的工作。您可以從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 。來自各種 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ù)依賴關(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 文件并打印其名稱。即使對(duì)于小型網(wǎng)絡(luò)，筆記本內(nèi)部的渲染也很麻煩，您可以在單獨(dú)的瀏覽器窗口中打開 SVG 文件進(jìn)行渲染。

TREx graphing API 是可配置的，允許使用各種顏色和格式，并且可用的格式設(shè)置程序包含信息。例如，使用默認(rèn)的格式化程序，層根據(jù)其操作進(jìn)行著色，并按名稱、類型和分析的延遲進(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ì)算精度可以提高吞吐量，降低延遲和功耗。

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

report_card_perf_overview(plan)

在檢查每層類型的延遲視圖時(shí)，共有 12 個(gè)重新格式化節(jié)點(diǎn)，約占運(yùn)行時(shí)的 26.5% 。那是相當(dāng)多的。在優(yōu)化過程中，重新格式化節(jié)點(diǎn)會(huì)插入到引擎圖中，但也會(huì)插入這些節(jié)點(diǎn)以轉(zhuǎn)換精度。每個(gè)重新格式化層都有一個(gè)原點(diǎn)屬性，描述其存在的原因。

如果您看到太多的精度轉(zhuǎn)換，您應(yīng)該看看是否可以做些什么來減少這些轉(zhuǎn)換。在 TensorRT 8.2 中，您可以看到縮放圖層，而不是為 Q / DQ 操作重新格式化圖層。這是因?yàn)?TensorRT 8.2 和 8.4 中使用了不同的圖形優(yōu)化策略。

圖 5 每層類型視圖的計(jì)數(shù)和延遲， ResNet18 QAT

要想挖得更深，請(qǐng)轉(zhuǎn)到衣料單元中可用的發(fā)動(dòng)機(jī)衣料 API 。您可以看到，卷積和 Q / DQ 過濾機(jī)都標(biāo)記了一些潛在的問題。

卷積 linter 標(biāo)記 13 個(gè)具有 INT8 輸入和 FP32 輸出的卷積。理想情況下，如果卷積后面是 INT8 精度層，則希望卷積輸出 INT8 數(shù)據(jù)。 linter 建議在卷積之后添加量化操作。為什么這些卷積的輸出沒有量化？

圖 6 卷積過濾機(jī)的輸出，關(guān)于帶浮點(diǎn)輸出的 INT8 卷積的警告

仔細(xì)看看。要在引擎圖中查找卷積，請(qǐng)從 linter 表中復(fù)制卷積的名稱，并在圖形 SVG 瀏覽器選項(xiàng)卡中搜索它。結(jié)果表明，這些卷積涉及到殘差加法運(yùn)算。

在咨詢了 Q / DQ 層鋪設(shè)建議 之后，您可能會(huì)得出結(jié)論，您必須在 PyTorch 模型中的剩余連接中添加 Q / DQ 層。不幸的是， QAT 工具包無法自動(dòng)執(zhí)行此操作，您必須手動(dòng)干預(yù) PyTorch 模型代碼。有關(guān)更多信息，請(qǐng)參閱 TensorRT QAT 工具包（ resnet.py ）中的示例。

下面的代碼示例顯示了BasicBlock.forward方法，新的量化代碼以黃色突出顯示。

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

圖 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è)引擎的延遲，按層類型分組

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

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

為什么間隙層在量化時(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 引擎瀏覽器，簡要回顧了它的 API 和特性，并通過一個(gè)示例演示了 TREx 如何幫助優(yōu)化 TensorRT 引擎的性能。 TREx 可以在 TensorRT 的 GitHub 存儲(chǔ)庫中的 實(shí)驗(yàn)工具 目錄下找到。

關(guān)于作者

Neta Zmora 是一位高級(jí)深度學(xué)習(xí)軟件架構(gòu)師，致力于 DL 加速。在 2020 年加入 NVIDIA 之前，內(nèi)塔是英特爾人工智能實(shí)驗(yàn)室的研究工程師，負(fù)責(zé)開發(fā)深度神經(jīng)網(wǎng)絡(luò)壓縮方法。

審核編輯：郭婷