施一飛

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

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

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

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

1 深度學(xué)習(xí)訓(xùn)練和部署

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

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

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

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

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

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

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

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

而經(jīng)過測試，在Jetson TX2平臺上，用GPU運(yùn)行TensorRT引擎優(yōu)化后的TensorFlow模型，推理速度提高了近4倍?？梢?，合理有效的利用TensorRT，可以極大的優(yōu)化目前嵌入式或汽車GPU平臺的AI推理速度和響應(yīng)時(shí)間。

【參考文獻(xiàn)】

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

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