探讨TensorRT加速AI模型的简易方案 — 以图像超分为例
AI模型近年來被廣泛應用于圖像、視頻處理，并在超分、降噪、插幀等應用中展現了良好的效果。但由于圖像AI模型的計算量大，即便部署在GPU上，有時仍達不到理想的運行速度。為此，NVIDIA推出了TensorRT，成倍提高了AI模型的推理效率。本次LiveVideoStack線上分享邀請到了英偉達DevTech團隊技術負責人季光一起探討把模型運行到TensorRT的簡易方法，幫助GPU編程的初學者加速自己的AI模型。
文 / 季光
整理 / LiveVideoStack
01
關于NVIDIA GPU
首先介紹英偉達的GPU。上一代GPU架構是圖靈Turing，當前架構是安培Ampere。Ampere消費級型號都是30開頭，包括3090、3080、3070等；企業級型號用于數據中心，包括A100、A30、A10、A16等。由于企業級型號很多，所以簡單介紹一下這些型號的用途。
A100是芯片面積最大的GPU，適合做訓練；而A30的能力大約是A100的一半。但這兩個GPU的特點是它們都支持新的數據格式TF32，并且在Tensor Core上做矩陣乘法有很高的吞吐（見上圖表格中標綠處）。TF32在訓練時非常有用，可以部分替代FP32。另外A100/A30支持MIG，可在單一操作系統中動態切割成多GPU，也可兼用于推理。
A10是T4的替代者，它的特點是FP32/FP16吞吐很高，比較適合做推理。
A16比較獨特，這個卡上含有4個GPU，每個GPU上帶著1個NVENC和2個NVDEC引擎，它更適合做轉碼。
GeForce 3090是消費型號，它的GPU型號與企業級的有所不同，計算能力有所欠缺，例如它的FP16的矩陣乘算力是142 TFLOPS（FP16累加，精度有限）或71 TFLOPS（FP32累加）。相比之下，A10的FP32累加矩陣乘可達125 TFLOPS，比它高出很多 。因此無論是做訓練還是做推理，GeForce 3090在很多情況下都比不過企業型號。
02
GPU編程基礎
GPU算力的發揮要靠GPU上的程序運行出來，因此需要我們編寫GPU的程序。GPU編程又被稱作異構編程，與CPU編程有不一樣的地方。
對于CPU程序，程序和數據都放在主存（即內存）上，這是我們熟悉的方式。而上圖左邊則是GPU程序的運行方式。GPU有自己的存儲器，即顯存。要把程序運行在GPU上時，我們需要先把數據從主存拷到顯存上，然后啟動GPU程序進行計算；當計算完成時，需要把數據從顯存拷回主存。以上就是異構編程的思想。簡單來說就是將數據拷至異構的處理器上，啟動程序，最后將數據拷回。
上圖右邊是個比較完整的程序，演示了上述思想。程序用cudaMalloc分配出顯存上的變量a和b（由顯存指針dp_a和dp_b指向），用cudaMemcpy把a從主存拷貝到顯存上，然后啟動GPU程序。黃色高亮的這段GPU程序稱作CUDA kernel，它所使用的數據都來自顯存。計算完成后，cudaMemcpy把結果b拷回主存，最后cudaFree釋放起初分配的顯存。
掌握“數據拷到顯存-啟動GPU程序-數據拷回主存”這一思想是非常重要的。對于熟悉C++編程的人來說，調用相關函數比較簡單，但要寫出CUDA kernel還需要額外花功夫。我們特別希望在使用GPU時可以減輕編程負擔，通過API調用方式就讓程序在GPU上運行起來。這也是TensorRT這種GPU加速庫出現的原因。
03
GPU轉碼流水線中的TensorRT
前面示例代碼中的數據是單個浮點數，這是一種簡單場景。而更復雜的場景下，拷貝的數據可以是單張圖片或連續圖片。無論如何，在主存和顯存間拷貝數據是有代價的，在數據量大時會成為程序運行的瓶頸，我們需要盡可能地減少或者避免。
以視頻轉碼為例，如果輸入數據是編碼過的視頻碼流，可以利用GPU上的硬件解碼器解碼，把解出的圖片存放在顯存，再交給GPU程序處理。此外，GPU上還帶有硬件編碼器，可以將處理后的圖片進行編碼，輸出視頻碼流。在上述流程中，無論是解碼，還是數據的處理，還是最后的編碼，都可以使數據留在顯存上，這樣可以實現較高的運行效率。
04
用TensorRT加速AI模型推理
深度學習應用的開發分為兩個階段，訓練和推理。TensorRT用來加速推理。
TensorRT的加速原理大體在這幾個方面：
TensorRT可以自動選取最優kernel。同樣是矩陣乘法，在不同GPU架構上以及不同矩陣大小，最優的GPU kernel的實現方式不同，TensorRT可以把它優選出來。
TensorRT可以做計算圖優化，通過kernel融合，減少數據拷貝等手段，生成網絡的優化計算圖。
TensorRT支持fp16/int18，對數據進行精度轉換，充分利用硬件的低精度、高通量計算能力。
05
TensorRT的加速效果
我們通過一些例子來說明TensorRT的加速效果。
對于常見的ResNet50來說，運行于T4，fp32精度有1.4倍加速；fp16精度有6.4倍加速。可見fp16很有用，啟用fp16相較于fp32有了進一步的4.5倍加速。
對于比較知名的視頻超分網絡EDVR，運行于T4，fp32精度有1.1倍加速，這不是很明顯；但fp16精度有2.7倍加速，啟用fp16相較于fp32有了進一步的2.4倍加速。
可以看出不同模型的加速效果不同，一般來說卷積模型加速較為顯著，而含大量數據拷貝的模型加速效果一般，且fp16無明顯幫助。
06
快速上手TensorRT
TensorRT該怎么用呢？本質上就是把訓練框架上訓練好的模型遷移到TensorRT上。以下是三種方案：
1）通過框架內部集成TensorRT
TensorFlow集成了TF-TRT，PyTorch還有TRTorch，調用這些API就可以把模型（部分地）運行在TensorRT上。它們的使用方式都比較簡單，通過框架中的API就能運行，但是很多情況下沒有達到最佳效率。
2）比較硬核的方法是使用TensorRT C++/Python API自行構造網絡，用TensorRT的API將框架中的計算圖重新搭一遍。這種做法兼容性最強，效率最高，但難度也最高。對于這種方法，我們之前在GTC China做過兩次報告（TENSORRT: 加速深度學習推理部署，利用 TENSORRT 自由搭建高性能推理模型?https://on-demand-gtc.gputechconf.com/gtcnew/sessionview.php?sessionName=ch8306-tensorrt%3a+加速深度學習推理部署），有興趣的話可以看一看，其難點是需要了解TensorRT的layer都有哪些，以及從原始框架的OP（即操作）跟這些layer的對應關系。
3）今天推薦的方法是從現有框架導出模型（ONNX）再導入TensorRT。
它的優點是難度適中，效率尚可，可以算作捷徑。需要解決的問題是：如何從訓練框架導出ONNX，以及如何把ONNX導入TensorRT。
07
解決如何導出與如何導入
第0步：了解TensorRT編程的基本框架
上圖展示的代碼是TensorRT最基本的使用方法。
1.作為準備工作，先造了logger，又造了builder，從builder造出network，這些對所有TensorRT程序都是固定的。
2.接下來高亮的這一部分是通過TensorRT的API把計算圖重建起來，使TensorRT上的計算與訓練框架原始模型一模一樣。這段代碼可以非常長，比如上百行。
3.做完之后利用network可以構建TensorRT engine（build_cuda_engine），構建時間根據網絡大小有長有短，短的幾秒，長的可達幾分鐘甚至幾小時。
4.構建好engine后可以調用運行。而且engine可以保存到磁盤，在第二次運行的時候，不需要再次build，直接load就可以運行。
上圖中的d_input、d_output是前面提到的異構編程中的顯存地址。
高亮的這一部分可以非常復雜，但為了省事，我們使用ONNX Parser自動搭建網絡，讓這一部分自動完成。
所以基本流程是這樣：先從訓練框架導出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不能跑也沒關系，我們仍有辦法讓TensorRT導入。這一點會在下文舉例說明。
上圖可以看到PyTorch導出ONNX的示例代碼。其中的resnet50是一個PyTorch nn.Module對象；verbose設為True可讓ONNX用文本方式打出來，對調試很有用；opset可以設置最高到12，版本越高，支持OP數量越多。
第2步：用Parser將ONNX導入TensorRT
TensorRT官方開發包自帶可執行文件trtexec。它可以接受ONNX輸入，根據ONNX將TensorRT網絡搭建起來，構建engine，并保存成文件。這一系列動作通過圖中的命令就可以做到。
其實trtexec也可以自己編程來實現，不過一般來說沒有必要。
trtexec運行成功說明TensorRT用自有的層重建了等價于ONNX的計算圖，而且計算圖被順利構建成了engine。保存成文件的engine將來可以反復使用。
第3步：運行Engine
最后一個步驟比較簡單，就是加載engine文件，提供輸入數據，即可運行。C++和Python的示例代碼可以從這里找到。(https://github.com/NVIDIA/trt-samples-for-hackathon-cn)
注意一定要對比TensorRT與原框架的計算結果，算出兩者的相對誤差均值。理想情況下fp32的誤差在1e-6數量級，fp16的誤差在1e-3數量級。
另外，我們都很關心模型跑到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端調用TensorRT的execute函數，其實是把GPU程序放進任務隊列，放好了就返回了，并不等GPU程序執行完畢；而GPU程序的執行卻是異步的。在CPU上做一次GPU同步，就是讓CPU等待此前提交的GPU任務全部執行完。基于以上，我們就可以理解為什么取時間之前要做一次GPU同步。
這里有個問題：這個簡易方法在什么時候不準確？簡單的說，這個方法會有誤差，如果要統計的GPU程序運行時間較短，就很難得出準確結果。這種時候，用CUDA event才是終極解決方法。
08
導出ONNX：疑難問題
前面說得都是最順利的情況。我們看看對于導出ONNX，不順利的情況有哪些：
如果遇到ONNX不支持的操作，解決辦法是升級框架和ONNX導出工具，使用當前支持的最高opset。
但這樣可能還不夠，因為有些PyTorch官方的OP在ONNX中仍然沒有定義（或無法組合得到）。所以在導出時加上選項ONNX_FALLTHROUGH，即便沒有定義也可以導出。
如果遇到開發者自定義的OP，則需要確認為自定義的Function子類增加symbolic函數，從而為自定義OP取ONNX節點名。
(例子見這里：https://github.com/shining365/EDVR-TRT/blob/master/basicsr/models/ops/dcn/deform_conv.py#L114)
此外，用trtexec把ONNX導入TensorRT時可能會遇到報錯。一種常見的情況是不支持的OP，這個稍后再說。另一種情況是TensorRT Parser對ONNX網絡結構有特殊要求。具體地，我們看一個例子。
上圖中高亮的報錯信息是”Resize scales must be an initializer！”為了得到更豐富的信息方便調試，請運行trtexec時打開--verbose選項。從圖中可以看到，這個Resize節點有385，402，401這3個輸入。這3個數字并不是輸入的具體值，而是輸入變量的名字。我們需要進一步看看這3個變量都是怎么生成的。
請在導出ONNX時確保設置verbose=True，可得到文本描述的ONNX，見上圖。可以看到Resize節點在圖中最下方，它的3個輸入變量已被高亮出來，它們有各自的計算過程。由于ONNX本身是個計算圖，我們可以畫一張圖將這一部分更清楚地展現出來。
09
ONNX手術刀：Graph Surgeon
上圖是有關這個Resize的ONNX子圖。它的第三個參數變量401來自Concat操作，將3個變量Concat在一起：其中一個是Constant，另外兩個是Constant經過了Unsqueeze與Cast，做了數據類型的轉換。
前文報錯信息“Resize scales must be an initializer！”指的是Resize的第三個參數不能是變量，而必須是Constant，所以我們需要把藍色的這部分子圖轉換成一個Constant，變成右邊的樣子。一旦做到，TensorRT Parser就會正常運行下去。
這個轉換在理論上可以做到，原因是這部分子圖的葉子節點都是Constant，具體值都寫在里面，我們按計算圖手工做一下相關計算，得到結果后存放在新建的Constant節點里就可以了。實現它的工具是Graph Surgeon。
Graph Surgeon像手術刀一樣可以修改ONNX計算圖。上圖就是用Graph Surgeon完成計算圖轉換的代碼。
1.首先找到符合條件的Resize節點，其篩選條件就是它的第三個輸入變量應來自Concat節點。
2.然后我們對這個Concat的所有輸入參數建立一個while循環，一直往上走，直到找到Constant，并把Constant里面的值放進values中。這樣走完for循環后，所有要合并的值都已經存進values中。
3.最后新建Constant節點，用numpy的concatenate函數將值合并填入該節點，并為該節點連接好輸出。
Graph Surgeon的一個完整示例代碼見這里(https://github.com/NVIDIA/trt-samples-for-hackathon-cn/blob/master/python/app_onnx_custom.py)。隨著大家做Graph Surgeon的經驗積累，特殊情況處理的經驗會越來越豐富，你將會積累更多的節點處理方式，從而讓更多模型被TensorRT Parser正確解析。
10
遇到不支持的操作
當trtexec報告不支持的OP時，我們不得不編寫TensorRT Plugin。TensorRT Plugin是TensorRT功能的擴展，需要什么我們就可以寫什么，也可以說是“萬金油”。
編寫TensorRT Plugin的思想是套用模板在里面“填空”。最關鍵的那個“空”就是GPU上的計算程序。對于缺少CUDA編程經驗的用戶，可以盡量復用原來代碼，避免新寫CUDA kernel。
這里我們演示了如何把EDVR里面的Deformable Convolution包裝成TensorRT plugin（代碼在這里：https://github.com/shining365/EDVR-TRT/blob/master/trt_onnx/DeformConvPlugin.h）。對于這個PyTorch的例子來說，我們盡量保持原始代碼不變，原封不動地把相關代碼片段提取出來，并拷貝了原始代碼的編譯選項，使得CUDA代碼可順利編譯。
11
使用fp16/int8加速計算
如果模型已經成功地跑在了TensorRT上，可以考慮使用fp16/int8做進一步加速計算。TensorRT默認運行精度是fp32；TensorRT在Volta、Turing以及Ampere GPU上支持fp16/int8的加速計算。
使用fp16非常簡單，在構造engine時設置標志即可。這一點體現在trtexec上就是它有--fp16選項，加上它就設置了這個標志。
我們舉例說明fp16加速計算的重要意義。對于EDVR，用ONNX導出的模型，直接運行fp32加速比是0.9，比原始模型慢，但是打開fp16就有了1.8倍加速。fp16對精度的影響不是很大。
int8量化需要校正數據集，而且這種訓練后量化一般會損失精度。如果對此介意，可以考慮使用Quantization Aware Training，做訓練時量化。
12
發揮TensorRT的極致性能
前面講的是TensorRT的一般用法，當你成為TensorRT熟手之后需要考慮如何發揮TensorRT極致性能。
1）API搭建網絡
對于EDVR來說，我在TensorRT上用過兩種方式運行，一種是用ONNX導出，它的fp32和fp16精度下的加速比是0.9和1.8；另一種是API搭建，它的加速比是1.1和2.7。可以看出API搭建有一定收益。假如模型特別重要，就要考慮用API搭建。
2）優化熱點
通過Nsight Systems可以找到時間占用最多的操作，對它進行重點優化。
3）用Plugin手工融合所有可以融合的層
以上這些方面都做到的話，基本上就可以做到在TensorRT上的極致性能。
13
總結與建議
今天我們推薦的開發方法是用ONNX Parser導入模型。這里需要熟悉Graph Surgeon用法，針對各種特殊情況處理。有可能需要自定義Plugin，包裝現有CUDA代碼。我們推薦使用混合精度，特別是fp16用法簡單、效果不錯；int8有更好計算性能，但一般會有精度下降。如果想要進階，要試著使用API搭建網絡，并且編寫與優化CUDA kernel。
14
示例代碼
以上就是我分享的全部內容，謝謝。
直播回放：
https://www.livevideostack.cn/video/gary-ji/
- The cover?from?creativeboom.com
講師招募?LiveVideoStackCon 2021 北京站
LiveVideoStackCon 2021 北京站（9月3-4日）正在面向社會公開招募講師，歡迎通過?speaker@livevideostack.com?提交個人及議題資料，無論你的公司大小，title高低，老鳥還是菜鳥，只要你的內容對技術人有幫助，其他都是次要的，我們將會在24小時內給予反饋。
總結
以上是生活随笔為你收集整理的探讨TensorRT加速AI模型的简易方案 — 以图像超分为例的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： LibAOM与AV1的最新研发进展
- 下一篇： LiveVideoStack主编观察04