AI模型近年來被廣泛應用于圖像、視頻處理，并在超分、降噪、插幀等應用中展現(xiàn)了良好的效果。但由于圖像AI模型的計算量大，即便部署在GPU上，有時仍達不到理想的運行速度。為此，NVIDIA推出了TensorRT，成倍提高了AI模型的推理效率。本次LiveVideoStack線上分享邀請到了英偉達DevTech團隊技術負責人季光一起探討把模型運行到TensorRT的簡易方法，幫助GPU編程的初學者加速自己的AI模型。

文 / 季光

整理 / LiveVideoStack

01

關于NVIDIA GPU

首先介紹英偉達的GPU。上一代GPU架構是圖靈Turing，當前架構是安培Ampere。Ampere消費級型號都是30開頭，包括3090、3080、3070等；企業(yè)級型號用于數(shù)據(jù)中心，包括A100、A30、A10、A16等。由于企業(yè)級型號很多，所以簡單介紹一下這些型號的用途。

• A100是芯片面積最大的GPU，適合做訓練；而A30的能力大約是A100的一半。但這兩個GPU的特點是它們都支持新的數(shù)據(jù)格式TF32，并且在Tensor Core上做矩陣乘法有很高的吞吐（見上圖表格中標綠處）。TF32在訓練時非常有用，可以部分替代FP32。另外A100/A30支持MIG，可在單一操作系統(tǒng)中動態(tài)切割成多GPU，也可兼用于推理。

• A10是T4的替代者，它的特點是FP32/FP16吞吐很高，比較適合做推理。

• A16比較獨特，這個卡上含有4個GPU，每個GPU上帶著1個NVENC和2個NVDEC引擎，它更適合做轉(zhuǎn)碼。

• GeForce 3090是消費型號，它的GPU型號與企業(yè)級的有所不同，計算能力有所欠缺，例如它的FP16的矩陣乘算力是142 TFLOPS（FP16累加，精度有限）或71 TFLOPS（FP32累加）。相比之下，A10的FP32累加矩陣乘可達125 TFLOPS，比它高出很多 。因此無論是做訓練還是做推理，GeForce 3090在很多情況下都比不過企業(yè)型號。

02

GPU編程基礎

GPU算力的發(fā)揮要靠GPU上的程序運行出來，因此需要我們編寫GPU的程序。GPU編程又被稱作異構編程，與CPU編程有不一樣的地方。

對于CPU程序，程序和數(shù)據(jù)都放在主存（即內(nèi)存）上，這是我們熟悉的方式。而上圖左邊則是GPU程序的運行方式。GPU有自己的存儲器，即顯存。要把程序運行在GPU上時，我們需要先把數(shù)據(jù)從主存拷到顯存上，然后啟動GPU程序進行計算；當計算完成時，需要把數(shù)據(jù)從顯存拷回主存。以上就是異構編程的思想。簡單來說就是將數(shù)據(jù)拷至異構的處理器上，啟動程序，最后將數(shù)據(jù)拷回。

上圖右邊是個比較完整的程序，演示了上述思想。程序用cudaMalloc分配出顯存上的變量a和b（由顯存指針dp_a和dp_b指向），用cudaMemcpy把a從主存拷貝到顯存上，然后啟動GPU程序。黃色高亮的這段GPU程序稱作CUDA kernel，它所使用的數(shù)據(jù)都來自顯存。計算完成后，cudaMemcpy把結果b拷回主存，最后cudaFree釋放起初分配的顯存。

掌握“數(shù)據(jù)拷到顯存-啟動GPU程序-數(shù)據(jù)拷回主存”這一思想是非常重要的。對于熟悉C++編程的人來說，調(diào)用相關函數(shù)比較簡單，但要寫出CUDA kernel還需要額外花功夫。我們特別希望在使用GPU時可以減輕編程負擔，通過API調(diào)用方式就讓程序在GPU上運行起來。這也是TensorRT這種GPU加速庫出現(xiàn)的原因。

03

GPU轉(zhuǎn)碼流水線中的TensorRT

前面示例代碼中的數(shù)據(jù)是單個浮點數(shù)，這是一種簡單場景。而更復雜的場景下，拷貝的數(shù)據(jù)可以是單張圖片或連續(xù)圖片。無論如何，在主存和顯存間拷貝數(shù)據(jù)是有代價的，在數(shù)據(jù)量大時會成為程序運行的瓶頸，我們需要盡可能地減少或者避免。

以視頻轉(zhuǎn)碼為例，如果輸入數(shù)據(jù)是編碼過的視頻碼流，可以利用GPU上的硬件解碼器解碼，把解出的圖片存放在顯存，再交給GPU程序處理。此外，GPU上還帶有硬件編碼器，可以將處理后的圖片進行編碼，輸出視頻碼流。在上述流程中，無論是解碼，還是數(shù)據(jù)的處理，還是最后的編碼，都可以使數(shù)據(jù)留在顯存上，這樣可以實現(xiàn)較高的運行效率。

04

用TensorRT加速AI模型推理

深度學習應用的開發(fā)分為兩個階段，訓練和推理。TensorRT用來加速推理。

TensorRT的加速原理大體在這幾個方面：

• TensorRT可以自動選取最優(yōu)kernel。同樣是矩陣乘法，在不同GPU架構上以及不同矩陣大小，最優(yōu)的GPU kernel的實現(xiàn)方式不同，TensorRT可以把它優(yōu)選出來。

• TensorRT可以做計算圖優(yōu)化，通過kernel融合，減少數(shù)據(jù)拷貝等手段，生成網(wǎng)絡的優(yōu)化計算圖。

• TensorRT支持fp16/int18，對數(shù)據(jù)進行精度轉(zhuǎn)換，充分利用硬件的低精度、高通量計算能力。

05

TensorRT的加速效果

我們通過一些例子來說明TensorRT的加速效果。

• 對于常見的ResNet50來說，運行于T4，fp32精度有1.4倍加速；fp16精度有6.4倍加速。可見fp16很有用，啟用fp16相較于fp32有了進一步的4.5倍加速。

• 對于比較知名的視頻超分網(wǎng)絡EDVR，運行于T4，fp32精度有1.1倍加速，這不是很明顯；但fp16精度有2.7倍加速，啟用fp16相較于fp32有了進一步的2.4倍加速。

可以看出不同模型的加速效果不同，一般來說卷積模型加速較為顯著，而含大量數(shù)據(jù)拷貝的模型加速效果一般，且fp16無明顯幫助。

06

快速上手TensorRT

TensorRT該怎么用呢？本質(zhì)上就是把訓練框架上訓練好的模型遷移到TensorRT上。以下是三種方案：

1）通過框架內(nèi)部集成TensorRT

TensorFlow集成了TF-TRT，PyTorch還有TRTorch，調(diào)用這些API就可以把模型（部分地）運行在TensorRT上。它們的使用方式都比較簡單，通過框架中的API就能運行，但是很多情況下沒有達到最佳效率。

2）比較硬核的方法是使用TensorRT C++/Python API自行構造網(wǎng)絡，用TensorRT的API將框架中的計算圖重新搭一遍。這種做法兼容性最強，效率最高，但難度也最高。對于這種方法，我們之前在GTC China做過兩次報告（TENSORRT: 加速深度學習推理部署，利用 TENSORRT 自由搭建高性能推理模型?https://on-demand-gtc.gputechconf.com/gtcnew/sessionview.php?sessionName=ch8306-tensorrt%3a+加速深度學習推理部署），有興趣的話可以看一看，其難點是需要了解TensorRT的layer都有哪些，以及從原始框架的OP（即操作）跟這些layer的對應關系。

3）今天推薦的方法是從現(xiàn)有框架導出模型（ONNX）再導入TensorRT。

它的優(yōu)點是難度適中，效率尚可，可以算作捷徑。需要解決的問題是：如何從訓練框架導出ONNX，以及如何把ONNX導入TensorRT。

07

解決如何導出與如何導入

第0步：了解TensorRT編程的基本框架

上圖展示的代碼是TensorRT最基本的使用方法。

1.作為準備工作，先造了logger，又造了builder，從builder造出network，這些對所有TensorRT程序都是固定的。

2.接下來高亮的這一部分是通過TensorRT的API把計算圖重建起來，使TensorRT上的計算與訓練框架原始模型一模一樣。這段代碼可以非常長，比如上百行。

3.做完之后利用network可以構建TensorRT engine（build_cuda_engine），構建時間根據(jù)網(wǎng)絡大小有長有短，短的幾秒，長的可達幾分鐘甚至幾小時。

4.構建好engine后可以調(diào)用運行。而且engine可以保存到磁盤，在第二次運行的時候，不需要再次build，直接load就可以運行。

上圖中的d_input、d_output是前面提到的異構編程中的顯存地址。

高亮的這一部分可以非常復雜，但為了省事，我們使用ONNX Parser自動搭建網(wǎng)絡，讓這一部分自動完成。

所以基本流程是這樣：先從訓練框架導出ONNX，再用TensorRT自帶的工具trtexec把ONNX導入TensorRT構建成engine，最后編寫一個簡單的小程序加載并運行engine即可。

第1步：從框架中導出ONNX

ONNX是中立計算圖表示，PyTorch有TouchScript，TensorFlow有frozen graph，都是框架特有的對于計算圖持久化的辦法。ONNX是平臺中立的，理論上所有框架都可以支持的表示方法。

一般情況下導出的ONNX仍具備運行能力，但有時不能直接運行，而是需要補充ONNX Runtime。比如導出的ONNX中具有特殊的算符，例如Deformable Convolution，它不是ONNX標準OP，但通過擴展ONNX Runtime可以讓導出的ONNX跑起來。

但ONNX能不能運行并不是可被TensorRT順利導入的先決條件。也就是說，導出的ONNX不能跑也沒關系，我們?nèi)杂修k法讓TensorRT導入。這一點會在下文舉例說明。

上圖可以看到PyTorch導出ONNX的示例代碼。其中的resnet50是一個PyTorch nn.Module對象；verbose設為True可讓ONNX用文本方式打出來，對調(diào)試很有用；opset可以設置最高到12，版本越高，支持OP數(shù)量越多。

第2步：用Parser將ONNX導入TensorRT

TensorRT官方開發(fā)包自帶可執(zhí)行文件trtexec。它可以接受ONNX輸入，根據(jù)ONNX將TensorRT網(wǎng)絡搭建起來，構建engine，并保存成文件。這一系列動作通過圖中的命令就可以做到。

其實trtexec也可以自己編程來實現(xiàn)，不過一般來說沒有必要。

trtexec運行成功說明TensorRT用自有的層重建了等價于ONNX的計算圖，而且計算圖被順利構建成了engine。保存成文件的engine將來可以反復使用。

第3步：運行Engine

最后一個步驟比較簡單，就是加載engine文件，提供輸入數(shù)據(jù)，即可運行。C++和Python的示例代碼可以從這里找到。(https://github.com/NVIDIA/trt-samples-for-hackathon-cn)

注意一定要對比TensorRT與原框架的計算結果，算出兩者的相對誤差均值。理想情況下fp32的誤差在1e-6數(shù)量級，fp16的誤差在1e-3數(shù)量級。

另外，我們都很關心模型跑到TensorRT上有多少加速比。熟悉CUDA編程的朋友可以用CUDA event測量運行時間，但要注意stream要設置正確。另外還有一種較粗略的簡易方法：做一次GPU同步，然后取時間t0；啟動GPU程序；再做一次GPU同步，取t1，得t1-t0，這就是GPU程序的運行時間。

(示例代碼見這里：https://github.com/NVIDIA/trt-samples-for-hackathon-cn/blob/master/python/app_onnx_resnet50.py)

這里關鍵需要理解GPU同步的含義：GPU程序是從CPU啟動的，即在CPU端調(diào)用TensorRT的execute函數(shù)，其實是把GPU程序放進任務隊列，放好了就返回了，并不等GPU程序執(zhí)行完畢；而GPU程序的執(zhí)行卻是異步的。在CPU上做一次GPU同步，就是讓CPU等待此前提交的GPU任務全部執(zhí)行完。基于以上，我們就可以理解為什么取時間之前要做一次GPU同步。

這里有個問題：這個簡易方法在什么時候不準確？簡單的說，這個方法會有誤差，如果要統(tǒng)計的GPU程序運行時間較短，就很難得出準確結果。這種時候，用CUDA event才是終極解決方法。

08

導出ONNX：疑難問題

前面說得都是最順利的情況。我們看看對于導出ONNX，不順利的情況有哪些：

如果遇到ONNX不支持的操作，解決辦法是升級框架和ONNX導出工具，使用當前支持的最高opset。

但這樣可能還不夠，因為有些PyTorch官方的OP在ONNX中仍然沒有定義（或無法組合得到）。所以在導出時加上選項ONNX_FALLTHROUGH，即便沒有定義也可以導出。

如果遇到開發(fā)者自定義的OP，則需要確認為自定義的Function子類增加symbolic函數(shù)，從而為自定義OP取ONNX節(jié)點名。

(例子見這里：https://github.com/shining365/EDVR-TRT/blob/master/basicsr/models/ops/dcn/deform_conv.py#L114)

此外，用trtexec把ONNX導入TensorRT時可能會遇到報錯。一種常見的情況是不支持的OP，這個稍后再說。另一種情況是TensorRT Parser對ONNX網(wǎng)絡結構有特殊要求。具體地，我們看一個例子。

上圖中高亮的報錯信息是”Resize scales must be an initializer！”為了得到更豐富的信息方便調(diào)試，請運行trtexec時打開--verbose選項。從圖中可以看到，這個Resize節(jié)點有385，402，401這3個輸入。這3個數(shù)字并不是輸入的具體值，而是輸入變量的名字。我們需要進一步看看這3個變量都是怎么生成的。

請在導出ONNX時確保設置verbose=True，可得到文本描述的ONNX，見上圖。可以看到Resize節(jié)點在圖中最下方，它的3個輸入變量已被高亮出來，它們有各自的計算過程。由于ONNX本身是個計算圖，我們可以畫一張圖將這一部分更清楚地展現(xiàn)出來。

09

ONNX手術刀：Graph Surgeon

上圖是有關這個Resize的ONNX子圖。它的第三個參數(shù)變量401來自Concat操作，將3個變量Concat在一起：其中一個是Constant，另外兩個是Constant經(jīng)過了Unsqueeze與Cast，做了數(shù)據(jù)類型的轉(zhuǎn)換。

前文報錯信息“Resize scales must be an initializer！”指的是Resize的第三個參數(shù)不能是變量，而必須是Constant，所以我們需要把藍色的這部分子圖轉(zhuǎn)換成一個Constant，變成右邊的樣子。一旦做到，TensorRT Parser就會正常運行下去。

這個轉(zhuǎn)換在理論上可以做到，原因是這部分子圖的葉子節(jié)點都是Constant，具體值都寫在里面，我們按計算圖手工做一下相關計算，得到結果后存放在新建的Constant節(jié)點里就可以了。實現(xiàn)它的工具是Graph Surgeon。

Graph Surgeon像手術刀一樣可以修改ONNX計算圖。上圖就是用Graph Surgeon完成計算圖轉(zhuǎn)換的代碼。

1.首先找到符合條件的Resize節(jié)點，其篩選條件就是它的第三個輸入變量應來自Concat節(jié)點。

2.然后我們對這個Concat的所有輸入?yún)?shù)建立一個while循環(huán)，一直往上走，直到找到Constant，并把Constant里面的值放進values中。這樣走完for循環(huán)后，所有要合并的值都已經(jīng)存進values中。

3.最后新建Constant節(jié)點，用numpy的concatenate函數(shù)將值合并填入該節(jié)點，并為該節(jié)點連接好輸出。

Graph Surgeon的一個完整示例代碼見這里(https://github.com/NVIDIA/trt-samples-for-hackathon-cn/blob/master/python/app_onnx_custom.py)。隨著大家做Graph Surgeon的經(jīng)驗積累，特殊情況處理的經(jīng)驗會越來越豐富，你將會積累更多的節(jié)點處理方式，從而讓更多模型被TensorRT Parser正確解析。

10

遇到不支持的操作

當trtexec報告不支持的OP時，我們不得不編寫TensorRT Plugin。TensorRT Plugin是TensorRT功能的擴展，需要什么我們就可以寫什么，也可以說是“萬金油”。

編寫TensorRT Plugin的思想是套用模板在里面“填空”。最關鍵的那個“空”就是GPU上的計算程序。對于缺少CUDA編程經(jīng)驗的用戶，可以盡量復用原來代碼，避免新寫CUDA kernel。

這里我們演示了如何把EDVR里面的Deformable Convolution包裝成TensorRT plugin（代碼在這里：https://github.com/shining365/EDVR-TRT/blob/master/trt_onnx/DeformConvPlugin.h）。對于這個PyTorch的例子來說，我們盡量保持原始代碼不變，原封不動地把相關代碼片段提取出來，并拷貝了原始代碼的編譯選項，使得CUDA代碼可順利編譯。

11

使用fp16/int8加速計算

如果模型已經(jīng)成功地跑在了TensorRT上，可以考慮使用fp16/int8做進一步加速計算。TensorRT默認運行精度是fp32；TensorRT在Volta、Turing以及Ampere GPU上支持fp16/int8的加速計算。

使用fp16非常簡單，在構造engine時設置標志即可。這一點體現(xiàn)在trtexec上就是它有--fp16選項，加上它就設置了這個標志。

我們舉例說明fp16加速計算的重要意義。對于EDVR，用ONNX導出的模型，直接運行fp32加速比是0.9，比原始模型慢，但是打開fp16就有了1.8倍加速。fp16對精度的影響不是很大。

int8量化需要校正數(shù)據(jù)集，而且這種訓練后量化一般會損失精度。如果對此介意，可以考慮使用Quantization Aware Training，做訓練時量化。

12

發(fā)揮TensorRT的極致性能

前面講的是TensorRT的一般用法，當你成為TensorRT熟手之后需要考慮如何發(fā)揮TensorRT極致性能。

1）API搭建網(wǎng)絡

對于EDVR來說，我在TensorRT上用過兩種方式運行，一種是用ONNX導出，它的fp32和fp16精度下的加速比是0.9和1.8；另一種是API搭建，它的加速比是1.1和2.7。可以看出API搭建有一定收益。假如模型特別重要，就要考慮用API搭建。

2）優(yōu)化熱點

通過Nsight Systems可以找到時間占用最多的操作，對它進行重點優(yōu)化。

3）用Plugin手工融合所有可以融合的層

以上這些方面都做到的話，基本上就可以做到在TensorRT上的極致性能。

13

總結與建議

今天我們推薦的開發(fā)方法是用ONNX Parser導入模型。這里需要熟悉Graph Surgeon用法，針對各種特殊情況處理。有可能需要自定義Plugin，包裝現(xiàn)有CUDA代碼。我們推薦使用混合精度，特別是fp16用法簡單、效果不錯；int8有更好計算性能，但一般會有精度下降。如果想要進階，要試著使用API搭建網(wǎng)絡，并且編寫與優(yōu)化CUDA kernel。

14

示例代碼

以上就是我分享的全部內(nèi)容，謝謝。

直播回放：

https://www.livevideostack.cn/video/gary-ji/

- The cover?from?creativeboom.com

講師招募?LiveVideoStackCon 2021 北京站

LiveVideoStackCon 2021 北京站（9月3-4日）正在面向社會公開招募講師，歡迎通過?speaker@livevideostack.com?提交個人及議題資料，無論你的公司大小，title高低，老鳥還是菜鳥，只要你的內(nèi)容對技術人有幫助，其他都是次要的，我們將會在24小時內(nèi)給予反饋。

總結

以上是生活随笔為你收集整理的探讨TensorRT加速AI模型的简易方案 — 以图像超分为例的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。