如何使用框架訓(xùn)練網(wǎng)絡(luò)加速深度學(xué)習(xí)推理

從 TensorRT 7.0 開始， Universal Framework Format（ UFF ）被棄用。在本文中，您將學(xué)習(xí)如何使用新的 TensorFlow -ONNX- TensorRT 工作流部署經(jīng)過 TensorFlow 培訓(xùn)的深度學(xué)習(xí)模型。圖 1 顯示了 TensorRT 的高級工作流。

圖 1 。 TensorRT 是一種推理加速器。

首先，使用任何框架訓(xùn)練網(wǎng)絡(luò)。網(wǎng)絡(luò)訓(xùn)練后，批量大小和精度是固定的（精度為 FP32 、 FP16 或 INT8 ）。訓(xùn)練好的模型被傳遞給 TensorRT 優(yōu)化器，優(yōu)化器輸出一個優(yōu)化的運(yùn)行時（也稱為計(jì)劃）。。 plan 文件是 TensorRT 引擎的序列化文件格式。計(jì)劃文件需要反序列化才能使用 TensorRT 運(yùn)行時運(yùn)行推斷。

要優(yōu)化在 TensorFlow 中實(shí)現(xiàn)的模型，只需將模型轉(zhuǎn)換為 ONNX 格式，并使用 TensorRT 中的 ONNX 解析器解析模型并構(gòu)建 TensorRT 引擎。圖 2 顯示了高級 ONNX 工作流。

圖 2 。 ONNX 工作流。

在本文中，我們將討論如何使用 ONNX 工作流創(chuàng)建一個 TensorRT 引擎，以及如何從 TensorRT 引擎運(yùn)行推理。更具體地說，我們演示了從 Keras 或 TensorFlow 中的模型到 ONNX 的端到端推理，以及使用 ResNet-50 、語義分段和 U-Net 網(wǎng)絡(luò)的 TensorRT 引擎。最后，我們將解釋如何在其他網(wǎng)絡(luò)上使用此工作流。

下載 TensorFlow -onnx- TensorRT 后 – 代碼 tar 。 gz 文件，您還應(yīng)該從 Cityscapes dataset scripts repo 下載 labels.py ，并將其與其他腳本放在同一個文件夾中。

ONNX 概述

ONNX 是機(jī)器學(xué)習(xí)和深度學(xué)習(xí)模型的開放格式。它允許您將不同框架（如 TensorFlow 、 PyTorch 、 MATLAB 、 Caffe 和 Keras ）的深度學(xué)習(xí)和機(jī)器學(xué)習(xí)模型轉(zhuǎn)換為單一格式。

它定義了一組通用的運(yùn)算符、深入學(xué)習(xí)的通用構(gòu)建塊集和通用文件格式。它提供計(jì)算圖的定義以及內(nèi)置運(yùn)算符。可能有一個或多個輸入或輸出的 ONNX 節(jié)點(diǎn)列表形成一個無環(huán)圖。

ResNet ONNX 工作流示例

在這個例子中，我們展示了如何在兩個不同的網(wǎng)絡(luò)上使用 ONNX 工作流并創(chuàng)建一個 TensorRT 引擎。第一個網(wǎng)絡(luò)是 ResNet-50 。

工作流包括以下步驟：

將 TensorFlow / Keras 模型轉(zhuǎn)換為。 pb 文件。

將。 pb 文件轉(zhuǎn)換為 ONNX 格式。

創(chuàng)建 TensorRT 引擎。

從 TensorRT 引擎運(yùn)行推斷。

將模型轉(zhuǎn)換為。 pb

第一步是將模型轉(zhuǎn)換為。 pb 文件。以下代碼示例將 ResNet-50 模型轉(zhuǎn)換為。 pb 文件：

除了 Keras ，您還可以從以下位置下載 ResNet-50 ：

深度學(xué)習(xí)示例 GitHub 存儲庫：提供最新的深度學(xué)習(xí)示例網(wǎng)絡(luò)。您還可以看到 ResNet-50 分支，它包含一個腳本和方法來訓(xùn)練 ResNet-50v1 。 5 模型。

NVIDIA NGC 型號 ：它有預(yù)訓(xùn)練模型的檢查點(diǎn)列表。例如，在 ResNet-50v1 。 5 上搜索 TensorFlow ，并從 Download 頁面獲取最新的檢查點(diǎn)。

將。 pb 文件轉(zhuǎn)換為 ONNX

第二步是將。 pb 模型轉(zhuǎn)換為 ONNX 格式。為此，首先安裝 tf2onnx 。

安裝 tf2onnx 后，有兩種方法可以將模型從。 pb 文件轉(zhuǎn)換為 ONNX 格式。第二種方法是使用命令行。運(yùn)行以下命令：

從 ONNX 創(chuàng)建 TensorRT 引擎

要從 ONNX 文件創(chuàng)建 TensorRT 引擎，請運(yùn)行以下命令：

此代碼應(yīng)保存在引擎。 py 文件，稍后在文章中使用。

此代碼示例包含以下變量：

最大工作區(qū)大?。?在執(zhí)行時 ICudaEngine 可以使用的最大 GPU 臨時內(nèi)存。

構(gòu)建器創(chuàng)建一個空網(wǎng)絡(luò)（ builder.create_network（） ）， ONNX 解析器將 ONNX 文件解析到網(wǎng)絡(luò)（ parser.parse（model.read（）） ）。您可以為網(wǎng)絡(luò)（ network.get_input（0）.shape = shape ）設(shè)置輸入形狀，然后生成器將創(chuàng)建引擎（ engine = builder.build_cuda_engine（network） ）。要創(chuàng)建引擎，請運(yùn)行以下代碼示例：

在這個代碼示例中，首先從 ONNX 模型獲取輸入形狀。接下來，創(chuàng)建引擎，然后將引擎保存在。 plan 文件中。

運(yùn)行來自 TensorRT 引擎的推理：

TensorRT 引擎在以下工作流中運(yùn)行推理：

為 GPU 中的輸入和輸出分配緩沖區(qū)。

將數(shù)據(jù)從主機(jī)復(fù)制到 GPU 中分配的輸入緩沖區(qū)。

在 GPU 中運(yùn)行推理。

將結(jié)果從 GPU 復(fù)制到主機(jī)。

根據(jù)需要重塑結(jié)果。

下面的代碼示例詳細(xì)解釋了這些步驟。此代碼應(yīng)保存在推理。 py 文件，稍后將在本文中使用。

為第一個輸入行和輸出行確定兩個維度。您可以在主機(jī)（ h_input_1 、 h_output ）中創(chuàng)建頁鎖定內(nèi)存緩沖區(qū)。然后，為輸入和輸出分配與主機(jī)輸入和輸出相同大小的設(shè)備內(nèi)存（ d_input_1 ， d_output ）。下一步是創(chuàng)建 CUDA 流，用于在設(shè)備和主機(jī)分配的內(nèi)存之間復(fù)制數(shù)據(jù)。

在這個代碼示例中，在 do_inference 函數(shù)中，第一步是使用 load_images_to_buffer 函數(shù)將圖像加載到主機(jī)中的緩沖區(qū)。然后將輸入數(shù)據(jù)傳輸?shù)?GPU （ cuda.memcpy_htod_async（d_input_1， h_input_1， stream） ），并使用 context.execute 運(yùn)行推理。最后將結(jié)果從 GPU 復(fù)制到主機(jī)（ cuda.memcpy_dtoh_async（h_output， d_output， stream） ）。

ONNX 工作流語義分割實(shí)例

在本文 基于 TensorRT 3 的自主車輛快速 INT8 推理 中，作者介紹了一個語義分割模型的 UFF 工作流過程。

在本文中，您將使用類似的網(wǎng)絡(luò)來運(yùn)行 ONNX 工作流來進(jìn)行語義分段。該網(wǎng)絡(luò)由一個基于 VGG16 的編碼器和三個使用反褶積層實(shí)現(xiàn)的上采樣層組成。網(wǎng)絡(luò)在 城市景觀數(shù)據(jù)集 上經(jīng)過大約 40000 次迭代訓(xùn)練

有多種方法可以將 TensorFlow 模型轉(zhuǎn)換為 ONNX 文件。一種方法是 ResNet50 部分中解釋的方法。 Keras 也有自己的 Keras 到 ONNX 文件轉(zhuǎn)換器。有時， TensorFlow -to-ONNX 不支持某些層，但 Keras-to-ONNX 轉(zhuǎn)換器支持這些層。根據(jù) Keras 框架和使用的層類型，您可能需要在轉(zhuǎn)換器之間進(jìn)行選擇。

在下面的代碼示例中，使用 Keras-to-ONNX 轉(zhuǎn)換器將 Keras 模型直接轉(zhuǎn)換為 ONNX 。下載預(yù)先訓(xùn)練的語義分段文件 semantic_segmentation.hdf5 。

圖 3 顯示了網(wǎng)絡(luò)的體系結(jié)構(gòu)。

圖 3 ?；?VGG16 的語義分割模型。

與前面的示例一樣，使用下面的代碼示例創(chuàng)建語義分段引擎。

要測試模型的輸出，請使用 城市景觀數(shù)據(jù)集 。要使用城市景觀，必須具有以下功能： sub_mean_chw 和 color_map 。這些函數(shù)也用于 post ， 基于 TensorRT 3 的自主車輛快速 INT8 推理 。

在下面的代碼示例中， sub_mean_chw 用于從圖像中減去平均值作為預(yù)處理步驟， color_map 是從類 ID 到顏色的映射。后者用于可視化。

下面的代碼示例是上一個示例的其余代碼。必須先運(yùn)行上一個塊，因?yàn)樾枰x的函數(shù)。使用這個例子比較 Keras 模型和 TensorRT 引擎 semantic 。 plan 文件的輸出，然后可視化這兩個輸出。根據(jù)需要替換占位符 /path/to/semantic_segmentation.hdf5 和 input_file_path 。

圖 4 顯示了實(shí)際圖像和實(shí)際情況，以及 Keras 的輸出與 TensorRT 引擎的輸出的對比。如您所見， TensorRT 發(fā)動機(jī)的輸出與 Keras 的類似。

圖 4a 原始圖像 。

圖 4b 地面真相標(biāo)簽。

圖 4c 。 TensorRT 的輸出。

圖 4d ： Keras 的輸出。

在其他網(wǎng)絡(luò)上試試

現(xiàn)在您可以在其他網(wǎng)絡(luò)上嘗試 ONNX 工作流。有關(guān)分段網(wǎng)絡(luò)的好例子的更多信息，請參閱 GitHub 上的 具有預(yù)訓(xùn)練主干的分割模型 。

作為一個例子，我們用一個 ONNX 網(wǎng)絡(luò)來說明如何使用。本例中的網(wǎng)絡(luò)是來自 segmentation_models 庫的 U-Net 。在這里，我們只加載模型，而沒有對其進(jìn)行訓(xùn)練。您可能需要在首選數(shù)據(jù)集上訓(xùn)練這些模型。

關(guān)于這些網(wǎng)絡(luò)的一個重要點(diǎn)是，當(dāng)您加載這些網(wǎng)絡(luò)時，它們的輸入層大小如下所示：（ None ， None ， None ， 3 ）。要創(chuàng)建一個 TensorRT 引擎，您需要一個輸入大小已知的 ONNX 文件。在將此模型轉(zhuǎn)換為 ONNX 之前，請通過為其輸入指定大小來更改網(wǎng)絡(luò)，然后將其轉(zhuǎn)換為 ONNX 格式。

例如，從這個庫（ segmentation _ models ）加載 U-Net 網(wǎng)絡(luò)并為其輸入指定大?。?244 、 244 、 3 ）。在為推理創(chuàng)建了 TensorRT 引擎之后，做一個與語義分段類似的轉(zhuǎn)換。根據(jù)應(yīng)用程序和數(shù)據(jù)集的不同，可能需要使用不同的顏色映射。

我們之前提到的另一種下載方式是從 vz6 下載。它有一個預(yù)先訓(xùn)練模型的檢查點(diǎn)列表。例如，您可以在 TensorFlow 中搜索 UNet ，然后轉(zhuǎn)到 Download 頁面以獲取最新的檢查點(diǎn)。

總結(jié)

在這篇文章中，我們解釋了如何使用 TensorFlow-to-ONNX-to-TensorRT 工作流來部署深度學(xué)習(xí)應(yīng)用程序，并給出了幾個示例。第一個例子是 ResNet-50 上的 ONNX- TensorRT ，第二個例子是在 Cityscapes 數(shù)據(jù)集上訓(xùn)練的基于 英偉達(dá)數(shù)據(jù)中心深度學(xué)習(xí)產(chǎn)品性能 的語義分割。

關(guān)于作者

Houman 是 NVIDIA 的高級深度學(xué)習(xí)軟件工程師。他一直致力于開發(fā)和生產(chǎn) NVIDIA 在自動駕駛車輛中的深度學(xué)習(xí)解決方案，提高 DNN 的推理速度、精度和功耗，并實(shí)施和試驗(yàn)改進(jìn) NVIDIA 汽車 DNN 的新思想。他在渥太華大學(xué)獲得計(jì)算機(jī)科學(xué)博士學(xué)位，專注于機(jī)器學(xué)習(xí)

About Yu-Te Cheng

Yu-Te Cheng 是 NVIDIA 自主駕駛組高級深度學(xué)習(xí)軟件工程師，從事自駕領(lǐng)域的各種感知任務(wù)的神經(jīng)結(jié)構(gòu)搜索和 DNN 模型訓(xùn)練、壓縮和部署，包括目標(biāo)檢測、分割、路徑軌跡生成等。他于 2016 年獲得卡內(nèi)基梅隆大學(xué)機(jī)器人學(xué)碩士學(xué)位。

About Josh Park

Josh Park 是 NVIDIA 的汽車解決方案架構(gòu)師經(jīng)理。到目前為止，他一直在研究使用 DL 框架的深度學(xué)習(xí)解決方案，例如在 multi-GPUs /多節(jié)點(diǎn)服務(wù)器和嵌入式系統(tǒng)上的 TensorFlow 。此外，他一直在評估和改進(jìn)各種 GPUs + x86 _ 64 / aarch64 的訓(xùn)練和推理性能。他在韓國大學(xué)獲得理學(xué)學(xué)士和碩士學(xué)位，并在德克薩斯農(nóng)工大學(xué)獲得計(jì)算機(jī)科學(xué)博士學(xué)位

審核編輯：郭婷