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

施一飛

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

【關鍵詞】人工智能；TensorRT；TensorFlow；神經網絡

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

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

1 深度學習訓練和部署

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

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

2 用TensorRT優化訓練后神經網絡

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

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

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

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

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

而經過測試，在Jetson TX2平臺上，用GPU運行TensorRT引擎優化后的TensorFlow模型，推理速度提高了近4倍。可見，合理有效的利用TensorRT，可以極大的優化目前嵌入式或汽車GPU平臺的AI推理速度和響應時間。

【參考文獻】

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

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