10倍加速！愛奇藝超分辨模型加速實踐

隨着終端播放設備的升級，觀衆對於視頻的品質需求也逐步提升。需求從最開始的高清過渡到4K，最近8K也有開始流行的趨勢。除了對於分辨率提升的需求之外，視頻在採集的過程中，也難免引入一些瑕疵，如物體運動過快導致的模糊，壓縮算法導致的畫質降低，拍攝/燈光等參數設置不佳導致的細節缺失，噪點增加等。經典插幀算法一般採用插值等算法，雖然速度很快，但細節豐富的圖片放大之後都會比較模糊，去噪更是困難。深度學習方法的引入，因爲其龐大的參數空間，很好的擬合了畫質的降噪過程，從而在提升分辨率的時候可以提供更多的細節，實現畫質和分辨率的雙重提升。

但深度學習模型，相比傳統方法，其運行時間大幅的提升，單個視頻的處理可能要達到數小時或者數天，難以滿足對海量視頻進行生產的需求。本文在這種背景下，介紹了愛奇藝在視頻4K 超分模型上進行的優化加速和生產落地實踐，將4K超分模型的性能在GPU 上提升了10倍。

複雜模型的部署挑戰

1.1 Nvidia在模型加速上提供的方法

Nvidia在Volta架構後，引入了tensorcore，一種domain specific accelerator [DSA]，用於對深度學習中常見的矩陣運算做加速處理。該加速器取得了極大的成功，大幅的降低了大型深度學習模型的推理時間。但tensor core作爲DSA也繼承了其普遍缺點，它的編程邏輯特別複雜，普通的程序員在用底層暴露API的方式幾乎難以使得tensor core達到它的理論運行速度。爲了降低編程的門檻，Nvidia構建了TensorRT[1]框架，將若干針對特定input/output tensor shape，用手工打磨彙編的方式，得到對應最優tensorcore性能的kernel，並通過抽象的接口暴露給TensorRT上層。

TensorRT在模型編譯時，會針對當前模型的狀況，依次在各個合適的內核上運行，最終挑選出耗時最小的內核，固定在最終編譯的TensorRT engine中。Engine即爲最終的部署binary，TensorRT推理時會加載engine，提取對應的內核名稱，以及對應的啓動參數，按照模型的推理順序，依次啓動內核從而完成推理過程。

雖然TensorRT方便了模型藉助tensorcore得到極大加速，但是由於tensor core本身的複雜性，TensorRT在對外暴露的接口較少，且核心算子實現目前還是閉源，進行模型深度優化時使用方式還是限制較多。舉例來說，目前TensorRT內部算子都是以NCHW進行開發的且僅支持NCHW的tensor輸入，但tensor core底層又需要以NHWC進行輸入，中間會進行多次tensor reshape 而降低效率。

1.2 愛奇藝在複雜視頻推理模型優化上的實踐

爲了進一步提高模型推理性能，愛奇藝對TensorRT底層機制做了詳細的解析。通過本文，您將得到如下的知識點：

a. 如何對複雜模型推理進行 TensorRT的格式轉換。

b. TensorRT的int8量化推理內部機制，以及如何更好的提升視頻推理中int8量化模型的推理精度。

複雜模型TensorRT的轉換方法

對於TensorRT模型部署來說，相信用過的人都碰到一個很頭疼的問題，即某個算子不支持，或者對於torch模型來說，很多算子是需要開發者使用CUDA來自定義實現的。其實這個不是TensorRT一家的問題，對於TensorRT立志成爲的通用深度學習編譯器來說，深度學習模型和框架迭代非常快速，各種模型的計算需求層出不窮，想要歸一化成爲一個通用的IR表示是非常的困難，更不用說將模型推向性能的極致。

針對不支持op或者自定義CUDA kernel的處理方法，我們實踐中通常用二字口訣來解決，《拆》，《合》。

2.1 模型《拆》解

就一般意義來說，模型的推理過程，其實就是一個完整計算圖的重放過程。既然是計算圖，我們將其拆解爲子圖，並橋接對應的輸入輸出，那麼對於其計算結果來說，應該是沒有影響的。

所以在這裏，我們使用了一個技巧，將可以export爲onnx，並正常轉換爲TensorRT的子圖獨立出來，編譯爲TensorRT的engine。然後在一個獨立的執行文件中，將子圖的engine依次replay，中間原始模型未轉換的op，我們用CUDA kernel來進行橋接。

圖1爲對於EDVR[2]具體的拆解部分。對於EDVR來說，主要是DCN自定義op和pixelshuffle無法正常轉換onnx以及TensorRT，故將對應模塊排除在外。將對於指定模塊使用torch的nn.moduel重新定義一個新的class，即可將對應的分塊用onnx給export出來。

圖1 EDVR中對於原始計算圖的分割

2.2 算子融《合》

以EDVR模型爲例，其中有一個無法轉換的op是Pixelshuffle。該OP是超分網絡中一種常見的upsample操作，用於在網絡的末尾處將相關feature的size給提升到目標大小。因爲無法直接轉換到TensorRT的engine中，所以需要將其獨立出來實現爲自定義的 CUDA kernel，作爲前後卷積部分的橋接單元。

但直接這樣操作對加速是不友好的。在前面提到過，TensorRT中tensor core的輸入是嚴格NHWC的，且TensorRT本身是遵守NCHW。因此TensorRT對於單個engine來說，輸入會有一個NCHW至NHWC的轉換，在一系列操作之後準備輸出之時，又會有一個NHWC至NCHW的轉換。這也就意味着我們的橋接op方式會觸發三個kernel操作，而由於超分的像素尺寸特別大，三個kernel各自的運行時間也較長。

圖2 pixelshuffle在TensorRT中的融合

圖2中藍色虛線框圖部分即爲 pixelshuffle 計算部分，中間表示爲 TensorRT的原始的算子轉換和 pixelshuffle 橋接的計算圖，我們可以看到三個模塊間進行了多次NCHW 和 NHWC 的reshape，很明顯這三個kernel其實是可以被合到一起的。因爲從卷積A的結束到卷積B的開始，中間的像素只是按照一定的規律進行了三次重排，中間不涉及到任何的計算，且完全線性。

融《合》原理雖很簡單，但實現上卻很有技巧。由於TensorRT 閉源，如何消除重排過程，在不改動TensorRT本身的情況下不太可能。因爲在推理代碼中，前後的連續卷積的橋接其實就只有pixelshuffle一個kernel而已，推理代碼並不知道TensorRT內部倆個“冗餘”重排內核的存在。

所以在這裏，又要進行更爲細緻的拆解。將TensorRT拆解爲執行文件可以看到的一個個內核，而不僅僅是一個黑盒的engine二進制文件。在這裏，我們不詳述具體的拆解過程，有興趣的讀者可以自行搜索相關CUDA hook方法來得到類似的CUDA kernel重放方法。

簡單的來說，我們用“錄音機”將TensorRT的運行軌跡給錄製成爲了“磁帶”。並且將磁帶按照歌單的順序把磁帶剪成了磁條，並且可以按照新的順序來重新播放“音樂”。這樣原來被隱藏的重排內核就暴露在了執行文件面前，按照其邏輯，我們手工優化了這一內核，從而代替了三個內核的運行。

TensorRT的int8推理

在Volta架構剛剛推出tensor core加速器的時候，Nvidia只支持了對fp16的支持。對於int8的完整支持，是到了後面Turing架構開始添加，int4/int1更是到了ampere架構才加入。Tensor core雖然是Nvidia應對一衆深度學習硬件加速器成功的反擊，但軟件支持如前所述，有着很大的使用限制。在量化方面，同樣如此。

TensorRT對於量化的支持要比其對tensor core的支持更早。大概從TensorRT4開始，就已經開始了對int8量化的支持[3]。最開始的量化是用從Pascal架構開始支持的dp4a指令來實現的，該指令可以將32bit運算並行化爲4個8bit運算，從而使得int8推理速度在當時架構上得到一個質的飛躍。

在TensorRT支持tensor core之後，也一直沿用着當時的框架，但只支持後訓練量化，採用KL散度逼近的方式，從全精度模型求解出對應的int8量化模型。這不是說TensorRT內部就不支持除後訓練以外的其他方法，但對於一般用戶來說，所能接觸的只有TensorRT暴露的API，API不支持後訓練量化以外的方法，那麼就意味着普通用戶與其他方法的絕緣。

3.1 TensorRT int8推理的內部機制

從TensorRT7開始，Nvidia開始將int8的量化過程以更精細的方式暴露出來。這個一方面也是由於torch/tensorflow開始支持了僞量化過程，另外onnx也提供了僞量化的op表示形式。但很遺憾的是TensorRT7對於這種全新int8的轉換方式支持還是有問題的，其中一個最大的問題就是卷積中的bias係數轉換的時候弄錯了，本應該乘的係數，變成了除，導致加入CNN中如果卷積有bias，那麼它的精度將大幅下降。【注：該問題已經在最新的TensorRT8中修復】

由於業務模型上線的壓力，不可能等到Nvidia出下一個版本來修復這個問題，於是我們又將視角投向了拆解。

爲了詳細的瞭解TensorRT int8的運算過程，我們對int8卷積內核做了反彙編，並對彙編代碼做了詳細的解讀。在瞭解了對應的彙編代碼之後，我們明白了其實int8運算過程中並不都是整形進行計算，中間是穿插着浮點運算的。

圖3 量化縮放過程

TensorRT在做模型轉換的時候，會將權值weight如圖3所示，以一個合適的縮放因子SW將原始浮點的範圍縮放到[-128,127]的int8整形範圍。同時在對模型進行finetune/calibration的時候，會對特定卷積層的輸入輸出總結出其對應的數據範圍，通過SI/SO從而能夠將原本也是浮點的輸入和輸出縮放到int8範圍。在engine二進制文件內部，權值weight就已經被固化成爲了int8的形式，並且輸入在進入卷積層之前就已經是int8的形式。這樣在卷積算子中的輸入/權值卷積乘加計算中，就都是int的形式，而由於int8的乘積非常容易產生溢出，卷積的乘加累加過程是以int32形式存在的。最終乘加的結果也是int32.乘加之後需要和bias進行進一步加法計算。但之前的int32結果其實是有輸入和權值對應的縮放因子SW/SI在裏面，所以爲了和bias的比例一致，需要將原來乘加結果的int32，先轉換爲浮點，然後依次除以SI和SW，再加上bias才能得到正確的結果，此時仍爲浮點類型，爲了以整形進行輸出，需要乘以SO纔得到最終輸出結果。

這裏以公式來表示即：

(IQ*WQ*SI/SW+B)*SO

該公式進而轉換爲

IQ*WQ*SO*SI/SW+B*SO

在這裏，TensorRT做了一個優化，它將SO*SI/SW合併爲一個新的係數，並直接將B*SO的結果存儲在engine二進制文件中，這樣卷積乘加後只再需要一次FMA操作即可獲得最終的結果。整個過程可以參加圖4。

圖4 TensorRT int8內核內部數值分佈情況

我們在TensorRT7的基礎上，抽取了內核的參數，並給其中的權值/bias進行重新賦值，來解決了原來實現中的問題。

3.2 進一步提升int8推理的精度

Int8雖然大幅的提升了推理的效率，同時QAT量化相比PTQ來說提升了量化的精度，但整體上int8相比全精度來說，仍然要有所下降。爲了進一步的提升int8模型的推理精度，我們採用了將TensorRT內核嵌入finetune過程和實時縮放因子計算這兩種方法。

TensorRT內核嵌入finetune過程

QAT finetune過程是一個僞量化的過程。對於pytorch來說，它僅僅是在輸入和輸出上做了一些事情，將權值/輸入/輸出按照圖x的縮放過程計算了一番，並縮放到整形範圍後，用round進行了一次截斷處理，隨後又使用相同的係數返回之前的浮點。但這一番操作引入的量化誤差就比較好的模擬了實際硬件中部署算子的操作。

雖然 pytorch 本身的量化是儘量的模擬實際量化推理計算的過程，但和實際推理計算相比，例如TensorRT的int8算子的計算結果還是有差異，這個差異對於最終的推理來說就是一個誤差的來源。

爲了消除這個誤差來源，我們將TensorRT內核嵌入pytorch 量化訓練的finetune過程中，確保量化訓練時使用的計算算子和 TensorRT 推理時的算子一致。簡單的說，就是將之前的“錄音機”又拿過來了，且因爲是訓練過程使用，內核的權值參數要計算前根據當前的數值進行更新。

縮放因子實時計算

在3.1闡述TensorRT內部int8計算機制的時候，提到過輸入/輸出的縮放因子是根據finetune數據集得到的一個經驗數值範圍計算得來的。這個過程存在另外一個誤差，即實際推理的時候，很有可能視頻幀的內容差異導致卷積生成feature map的數值範圍和finetune數據集中的範圍不一致。這樣繼續沿用老的數值，就會導致新的內容推理結果精度的下降。

爲了解決這個問題，我們引入了對於縮放因子，特別是輸出縮放因子的動態更新。一般來說，對於輸入（即原始幀），它的縮放因子都是可以固定的，而如果輸出縮放因子可以動態計算出來，就可以由級聯關係繼承下去變爲下層卷積的輸入縮放因子，從而整個網絡的各個縮放因子得到更新。

在這裏，我們又進一步拆解了TensorRT，給原始int8內核增加了一個新的彙編模塊，在這裏int8卷積的輸出不再是int8，而變爲float16類型。這樣的改變使得對於卷積算子來說就不再需要輸出的縮放因子的參與。在緊跟着的內核中，使用reduce方式計算出整體float16的最大值，從而確定出該輸出的縮放因子的數值，進而將該float16輸出縮放爲int8供下一層輸入使用。整體過程如圖5所示。

圖5 整體int8精度優化過程

性能提升結果

在整個EDVR部署優化的過程中，除了文章提到的《拆》《合》優化之外，同時也包含了其他的一些優化。如圖6所示中的，我們在優化的第二個步驟中，集中優化了DCN自定義op中的冗餘顯存訪問，從而大幅的提升了自定義算子自身的效率。在第三個步驟中，我們將一些算子如leaky進行融合，隨後獲取了大約150ms的收益。第四以及第五步驟中，我們集中對於一些中間態的格式轉換做了相互消減操作，使得在fp16精度上，EDVR達到了380ms的速度，最終int8的成功應用使得模型的推理效率進一步提升至1080p上180ms單幀的速度。