對使用TensorRT加速AI深度學習推斷效率的探索

施一飛

【摘 要】本文對使用TensorRT對目前流行的人工智能框架訓練的神經(jīng)網(wǎng)絡模型的優(yōu)化做了初步探索。目前基于人工智能的應用呈現(xiàn)爆炸式增長，圖像識別、語音識別，自然語言處理、圖像檢索等功能已經(jīng)成了人們?nèi)粘１仨毜墓ぞ?，隨之而來對人工智能的推理效率和響應速度則提出更高要求。本文結合英偉達（nVidia）的TensorRT 3引擎和面向無人智能化的Jetson TX2平臺，對以上問題的解決做了初步嘗試。

【關鍵詞】人工智能；TensorRT；TensorFlow；神經(jīng)網(wǎng)絡

作為人工智能浪潮中數(shù)字產(chǎn)品和服務的消費者，我們每天都會與多種人工智能相關的服務進行交互，包括語音識別，語言翻譯，圖像識別或視頻字幕生成等。在這些產(chǎn)品和服務的背后，神經(jīng)網(wǎng)絡計算每個功能請求的結果。這個步驟在人工智能領域通常被稱為“推理”，即新的數(shù)據(jù)通過傳遞給受過訓練的神經(jīng)網(wǎng)絡計算以產(chǎn)生結果。在過往的機器學習專業(yè)中，有時稱這個過程為“預測”或“計分”。

目前流行的做法是將神經(jīng)網(wǎng)絡部署在云端Web服務中，接收來自數(shù)千或數(shù)百萬用戶的并發(fā)請求，計算每個請求的推理結果并響應給用戶。為了提供良好的用戶體驗，在考慮包括網(wǎng)絡延遲，神經(jīng)網(wǎng)絡執(zhí)行消耗和基于生產(chǎn)環(huán)境的其他延遲的情形下，努力促成最少響應時間。業(yè)界主要通過GPU的多核心并發(fā)計算能力和浮點運算的優(yōu)勢，利用諸如Caffe、TensorFLow等深度學習框架加速深度神經(jīng)網(wǎng)絡的訓練，解決特定應用程序的效率要求。這些框架可以進行網(wǎng)絡設計，執(zhí)行模型評估和診斷以及用新數(shù)據(jù)重新訓練模型，為快速構建問題解決原型提供了很大的靈活性。

1 深度學習訓練和部署

用深度神經(jīng)網(wǎng)絡解決機器學習的監(jiān)督問題涉及兩個步驟：第一步，是使用GPU對海量標記數(shù)據(jù)進行深度神經(jīng)網(wǎng)絡訓練。這一步中，神經(jīng)網(wǎng)絡將學習數(shù)以百萬計的權重或參數(shù)，使其能夠在示例數(shù)據(jù)與正確的響應之間建立映射。因為目標函數(shù)相對于網(wǎng)絡權重已經(jīng)被最小化，訓練過程通常需要通過網(wǎng)絡實現(xiàn)前推回代。通常，模型訓練和精確度需要根據(jù)非訓練期間的數(shù)據(jù)進行驗證，以預估在實景的表現(xiàn)。第二步，是使用訓練好的模型根據(jù)新采集的數(shù)據(jù)來進行預測。該步驟中，訓練效果最好的模型被用于運行在生產(chǎn)環(huán)境中的應用程序，如數(shù)據(jù)中心，汽車或嵌入式平臺。對于某些應用，如自動（輔助）駕駛，推理預測是需要實時完成的，那么計算中對高吞吐量、響應時間的駕馭至關重要。

如果目標是使用訓練后的神經(jīng)網(wǎng)絡來感知其周圍的嵌入式設備，那么通過模型的前向推斷直接影響整體響應時間和設備消耗的功率。其中，要優(yōu)化的關鍵指標是功率效率，即每瓦推斷性能。每瓦性能也是數(shù)據(jù)中心運營效率最大化的關鍵指標。在這種情況下，需要盡量減少地理上和時間上大量不相關的請求所引起的延遲和消耗的能量，這往往也成了限制大批量計算能力的關鍵因素。

2 用TensorRT優(yōu)化訓練后神經(jīng)網(wǎng)絡

TensorRT通過優(yōu)化訓練后的神經(jīng)網(wǎng)絡來生成出于待部署狀態(tài)的推理引擎，能最大限度地提高GPU推理性能和功效，尤其可以為諸如圖像分類、分割和對象檢測之類的深度學習應用提供最大的推理吞吐量和效率。TensorRT運行時，對從P4和V100等數(shù)據(jù)中心GPU到自動駕駛和嵌入式平臺（如Drive PX2和Jetson TX2）的每個GPU平臺的依賴性最小。

使用TensorRT先后需要完成編編譯和部署兩個階段任務。在編譯階段，TensorRT對網(wǎng)絡配置進行優(yōu)化，生成優(yōu)化計劃，用于計算通過深度神經(jīng)網(wǎng)絡的正向推理。部署時采用駐守服務或用戶應用程序的形式接收輸入的批量數(shù)據(jù)，推理時根據(jù)輸入數(shù)據(jù)執(zhí)行優(yōu)化計劃，最后返回輸出數(shù)據(jù)（分類，對象檢測等）。

這里以TensorRT部署TensorFlow模型為例，說明如何在TensorRT中導入和優(yōu)化經(jīng)過訓練的TensorFlow神經(jīng)網(wǎng)絡并將其部署為TensorRT運行時引擎。需要完成兩個步驟：（1）導入并優(yōu)化訓練過的模型（這里需要使用TensorFlow框架的freeze_graph工具導出為凍結模型再導入），設置部署目標使用的GPU并發(fā)量、進度和內(nèi)存等優(yōu)化選項，生成推理引擎。以文件的形式存儲在磁盤上。（2）加載和反序列號把存好的推理引擎計劃文件，創(chuàng)建TensorRT引擎對象，使用它來進行目標部署平臺上的新數(shù)據(jù)推理。注意，這里我們可以不用安裝深度學習框架就可完成。下面是示例代碼：

代碼中UFF代表通用框架格式（Universal Framework format），它是運行優(yōu)化引擎前用于表示網(wǎng)絡模型的TensorRT內(nèi)部格式。函數(shù)from_tensorflow_frozen_model（）負責進行模型轉換，第一個參數(shù)是待導入的TensorFlow訓練模型。第二個參數(shù)是輸出圖層名稱。目前在TensorRT中支持一般的標準圖層類型，包括卷積層、縮放層、完全連接層等。上述代碼的輸出是TensorFlow模型的UFF圖形表示，它可以通過TensorRT進行分析。我們通過提供輸入層的名稱和維度（CHW格式）和輸出層的名稱來配置UFF解析器。接下來的編譯過程，需要在目標平臺進行編譯，這里我們使用的是nVidia Jetson TX2，在需要使用Jetson TX2完成編譯。在設置TensorRT的優(yōu)化參數(shù)時，設定推理并發(fā)量、內(nèi)存占用和目標進度：

參數(shù)指定并發(fā)量為1，使用1MB的暫存空間進行進度為FP32的推斷。函數(shù)的輸出就是TensorRT優(yōu)化的運行時引擎，可以用來進行推理。優(yōu)化的內(nèi)容包括圖層和張量垂直融合，內(nèi)核的自動調(diào)優(yōu)，張量動態(tài)記憶等。最后通過trt.utils.write_engine_to_file函數(shù)將引擎序列化為文件存儲。當需要執(zhí)行TensorRT的推斷時，加載序列化的模型，進行推理。

而經(jīng)過測試，在Jetson TX2平臺上，用GPU運行TensorRT引擎優(yōu)化后的TensorFlow模型，推理速度提高了近4倍?？梢姡侠碛行У睦肨ensorRT，可以極大的優(yōu)化目前嵌入式或汽車GPU平臺的AI推理速度和響應時間。

【參考文獻】

[1]NVIDIA CUDA Programming Guide. http：//docs.nvidia.com/cuda/pdf/CUDA_C_Programming_Guide.pdf.2013.

[2]面向GPU計算平臺的若干并行優(yōu)化關鍵技術研究[D]. 賈海鵬.中國海洋大學，2012.endprint