本文經(jīng)自動(dòng)駕駛之心公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

一. TensorRT是什么？

2016年Nvidia為自家GPU加速推理而提供的SDK，人們有時(shí)也把它叫做推理框架。

二. 為什么？

只有Nvidia最清楚自家GPU或DLA該如何優(yōu)化，所以TensorRT跑網(wǎng)絡(luò)的速度是最快的，比直接用Pytorch快N倍。

遙遙領(lǐng)先的TensorRT

三. 怎么做到的？

1. 搜索整個(gè)優(yōu)化空間

與Pytorch等其它訓(xùn)練框架最大區(qū)別是，TensorRT的網(wǎng)絡(luò)優(yōu)化算法是基于目標(biāo)GPU所做的推理性能優(yōu)化，而其它框架一方面需要綜合考慮訓(xùn)練和推理，更重要的是它們沒有在目標(biāo)GPU上做針對(duì)性的優(yōu)化。

TensorRT又是如何針對(duì)目標(biāo)GPU優(yōu)化的呢？

簡(jiǎn)單講就是在可能的設(shè)計(jì)空間中搜索出全局最優(yōu)解。

這個(gè)搜索空間有哪些變量呢？

比如CUDA架構(gòu)中的編程模型所對(duì)應(yīng)的，將Tensor劃分為多少個(gè)block？以及這些block如何組織到Grid中。

任務(wù)被劃分為多個(gè)Block

Block以Grid的方式組織起來(lái)

不同的組織層次以對(duì)應(yīng)不同的存儲(chǔ)體系結(jié)構(gòu)中的不同存儲(chǔ)器

再舉例，使用什么樣的指令完成計(jì)算，可能是FFMA、FMMA，可能是TensorCore指令...

更難的部分可能是Tensor數(shù)據(jù)流的調(diào)度，把他們放在local、share還是global memory呢？如何擺放呢？

這些變量組合在一起是一個(gè)巨大的搜索空間，可能你的CPU計(jì)算幾天也得不出個(gè)結(jié)果來(lái)。

但是，我們知道神經(jīng)網(wǎng)絡(luò)的計(jì)算是由一個(gè)個(gè)粒度更大的算子組成的，算子上面還有粒度更大的層結(jié)構(gòu)。我們也清楚地知道層與層之間相對(duì)獨(dú)立，也就是說可以針對(duì)每層計(jì)算優(yōu)化，最后把優(yōu)化后的層串在一起大概率就是網(wǎng)絡(luò)的全局最優(yōu)解。

于是，TensorRT預(yù)先寫了很多算子和層（CUDA Kernel）。當(dāng)然這些算子的輸入和輸出tensor是可以配置的，以適應(yīng)網(wǎng)絡(luò)輸入和輸出的不同以及GPU資源的不同。

部分優(yōu)化好的算子

搜索空間變小了，從原來(lái)的指令級(jí)別的搜索，上升到了算子級(jí)別的搜索。因?yàn)檫@些實(shí)現(xiàn)都是用CUDA kernel所寫，更準(zhǔn)確的說是Kernel級(jí)別的搜索了。

但是Tensor數(shù)據(jù)流的調(diào)度問題并沒有解決，這也是最關(guān)鍵和復(fù)雜的地方。我們應(yīng)該將輸入Tensor劃分為多少個(gè)Block呢？這些Blocks應(yīng)該分配給多少個(gè)線程呢？Tensor存儲(chǔ)在哪呢？local/share/global memory的哪些地方呢？中間計(jì)算結(jié)果存儲(chǔ)在哪里呢？

對(duì)于計(jì)算部分是可以通過模擬的方式（類似指令集仿真器）計(jì)算得到性能的，但是Tensor數(shù)據(jù)流在share/L2/Global Memory的流動(dòng)過程就很難通過仿真計(jì)算得到精確結(jié)果，因?yàn)橐荒M的數(shù)據(jù)量和線程數(shù)過大，何況要嘗試的可能性還很多，靠CPU仿真計(jì)算的思路就別想了。唯一辦法就是讓候選算子在目標(biāo)GPU上直接跑跑，統(tǒng)計(jì)出性能，最后通過比對(duì)選出最優(yōu)解。TensorRT把這個(gè)過程叫做Timing，TensorRT甚至可以將優(yōu)化的中間過程存儲(chǔ)下來(lái)供你分析，叫做timing caching（通過trtexec --timingCacheFile=<file>）。

Nvida GPU memory架構(gòu)

以上所描述的優(yōu)化過程可以叫做Hardware Aware Optimazation。

總結(jié)起來(lái)優(yōu)化器會(huì)重點(diǎn)分析：

• Type of hardware（Hardware capability...）
• Memory footprint（Share, Cache, Global...）
• Input and output shape
• Weight shapes
• Weight sparsity
• Level of quantization （so, reconsider memory)

而這些是Pytorch等框架不會(huì)去深入挖掘的，尤其是對(duì)存儲(chǔ)系統(tǒng)的優(yōu)化。

2. 強(qiáng)制選擇Kernel

由于Block之間線程的運(yùn)行順序是隨機(jī)的，CPU可能在向GDDR/HBM讀寫數(shù)據(jù)，甚至GPU的時(shí)鐘頻率也在隨負(fù)載的變化而變化，這導(dǎo)致了不同系統(tǒng)運(yùn)行環(huán)境下GPU的性能表現(xiàn)會(huì)有差異。這種差異也可能導(dǎo)致TensorRT Timing的最優(yōu)解不是實(shí)際推理時(shí)的最優(yōu)解，可能選擇了次優(yōu)的Kernel。

TensorRT提供了一個(gè)補(bǔ)救方法，就是強(qiáng)制指定選擇某個(gè)Kernel實(shí)現(xiàn)，如果你很確信它是最優(yōu)解的話。

TensorRT提供的API叫做AlgorithmSelector。

3. Plugin

當(dāng)然，你對(duì)自己設(shè)計(jì)的算子更有把握，可以自己寫Kernel，然后指定使用它。

不過，更多情況下，是因?yàn)榘l(fā)現(xiàn)TensorRT不支持某個(gè)算子，你才被迫去寫Kernel，畢竟CUDA編程不簡(jiǎn)單，何況性能還需要足夠好。

4. cuBLAS和cuDNN

TensorRT安裝指導(dǎo)要求你先安裝CUDA SDK和cuDNN。

CUDA SDK需要安裝是顯而易見的，因?yàn)門ensorRT所調(diào)用的Kernel需要NVCC編譯器來(lái)編譯成Nvidia GPU的匯編指令序列啊！

但是CUDA SDK中還有一個(gè)cuBLAS庫(kù)也是被TensorRT所依賴的，我們知道C++庫(kù)BLAS（Basic Linear Algebra Subprograms），它是針對(duì)CPU進(jìn)行的線性代數(shù)計(jì)算優(yōu)化，那么cuBLAS就是針對(duì)CUDA GPU開發(fā)的線性代數(shù)計(jì)算庫(kù)，它的底層當(dāng)然也就是用CUDA Kernel寫成的。典型的矩陣乘法算子就可以直接調(diào)用cuBLAS了。

cuBLAS開發(fā)的很早，應(yīng)該是CUDA生態(tài)最早的一批庫(kù)了吧，但是隨著深度學(xué)習(xí)的普及，Nvidia又在生態(tài)中加入了cuDNN庫(kù)，它的層次更高，封裝了到了網(wǎng)絡(luò)層，所以其實(shí)TensorRT也可以直接調(diào)用優(yōu)化好的cuDNN庫(kù)中的Kernel？是也不是。

TensorRT可以選擇所謂Tactic（策略）來(lái)決定是使用TensorRT寫的Kernel還是cuBLAS和cuDNN的。

5. Tactic

TensorRT的Tactic能決定很多優(yōu)化選項(xiàng)。

例如，每次timing某個(gè)算子時(shí)需要平均的運(yùn)行次數(shù)。缺省TensorRT會(huì)運(yùn)行四次，以降低不確定性帶來(lái)的誤差，但這個(gè)次數(shù)是可以修改的。

還可以決定上面提到的Kernel庫(kù)的選擇，Plugin的選擇，GPU時(shí)鐘頻率鎖定等。

6. 量化

TensorRT當(dāng)然具備網(wǎng)絡(luò)量化能力，提供了將全網(wǎng)都量化到int8的隱性量化方式，也提供了插入Q/DQ Layer的顯性量化方式。

混合量化是Nvidia做的很優(yōu)秀的地方，這對(duì)于高效利用計(jì)算資源起到了重要作用，不過，這個(gè)另外的話題，以后有機(jī)會(huì)再談。

7. 多應(yīng)用推理和多卡推理

其實(shí)這才是Nvidia強(qiáng)悍的地方，在友商都在談單卡性能時(shí)，其實(shí)多卡或多節(jié)點(diǎn)才是Nvidia的殺手锏

另外，對(duì)于單卡性能富余的情況下，可能希望有多個(gè)流并行推理，這個(gè)對(duì)于TensorRT來(lái)說也是必須支持的

四. TensorRT的內(nèi)核到底是什么？

答：根據(jù)網(wǎng)絡(luò)、輸入、輸出tensor、目標(biāo)GPU的資源，通過實(shí)際運(yùn)行，在候選Kernel庫(kù)中擇優(yōu)的一個(gè)Hardware Aware優(yōu)化器。

五. 編譯器

最后，如果非要套用編譯器前后端理論的話，上述談到的部分應(yīng)該屬于編譯器后端部分了，因?yàn)樗呀?jīng)和底層硬件息息相關(guān)了。只不過它邏輯上處于于NVCC這個(gè)實(shí)體編譯器的上層。而編譯器前端，也就是與硬件不相關(guān)的圖融合部分是也是在TensorRT的Builder內(nèi)完成的。

好了，如果你對(duì)AI編譯器還不了解，可以看下面這篇入門文章

https://zhuanlan.zhihu.com/p/632648673

最后送上兩幅圖，作為總結(jié)

TensorRT工具鏈

TensorRT后端優(yōu)化流程