量化深度學習模型在CUDA平臺的自動優化

深度學習已經成功地在各種任務中得到應用。模型的推斷速度在無人駕駛等實時的場景下尤為關鍵。網路量化是加速深度學習模型一種有效的方法。在量化的模型中，我們使用 int8、float16 等低精度的資料型別表示資料和模型引數。更低的資料頻寬不僅降低了模型推斷時間和記憶體及儲存上的需求，也降低了執行時的功耗。與此同時，通過採用合適的 calibration （校準）方案，我們能夠把模型精度損失降到最低。模型量化使得較大的模型能夠在GPU、CPU、移動端裝置等不同的硬體上部署。因此，模型量化在近年來受到大量的研究者和開發者的關注。

在此之前，量化模型中的運算子通常使用人工編寫的 microkernel 來實現，或者依賴cuDNN、TensorRT等以黑盒形式提供的商用軟體。通過彙編編寫高效能的 microkernel 是一項困難且耗費大量人力的工作；並且，這些 microkernel 通常無法在新出現的模型和裝置上達到較好的效能。

針對這個挑戰，TVM 使用一個完整的編譯器棧和基於機器學習的優化器來自動地生成高效的程式碼。通過在人工設計的搜尋空間內自動搜尋，TVM 能夠生成高效的 kernel 程式碼。在VGG、ResNet等標準的模型上，TVM 達到了和 TensorRT 近似的效能。而在ResNext、Deformable ConvNets等新出現的模型上，利用TVM的自動優化我們能夠十分容易地實現其他框架無法做到的效能提升。

我們將在本篇文章中介紹 TVM 在 CUDA 平臺上對 int8 量化深度學習模型的優化。

使用 TVM 實現量化 CUDA kernel

通過 tensorization 利用 intrinsic 指令

許多平臺提供了特定架構下的優化指令來處理一些特殊的計算。例如，x86 平臺的 SIMD 指令、CUDA 平臺的 dp4a 和 hfma 指令等。這些指令通常高度優化。利用這些 intrinsic 指令，我們能夠實現顯著的效能提升。

我們在 TVM 的 CUDA int8 運算子中使用dp4a 指令進行優化。dp4a 是 CUDA 平臺 Compute Capability 6.1 中的用於計算長度為 4 的 8 位整數向量點積，併產生32位計算結果的一條 intrinsic 指令。我們使用 dp4a 實現任意滿足長度為 4 的倍數的 int8 向量點積。而卷積、全連線等通常依賴點積運算，藉助高效的向量點積我們能夠更進一步地實現這些更高層次的運算子。

例如，2D卷積可以直接在 channel、width、height 維度累加進行計算。這是一個 dp4a 的典型應用。在TVM中，我們使用 tensorization（張量化）來支援外部指令的呼叫。這一步不需要修改計算過程的宣告，只需在 TVM 中的 schedule 階段呼叫 “tensorize” 這一個 schedule primitive，將累加部分替換為基於 dp4a 的實現。關於 tensorization 的細節可以檢視這篇教程。

資料佈局的重排

在做 tensorization的時候，我們可能需要對計算邏輯進行特殊的設計。在全連線運算子中，我們可以直接沿著陣列內層索引的順序依次累加。而卷積運算的實現更為複雜。在2d卷積中，我們希望沿著 channel 這一維度使用 dp4a 指令進行運算，因為大多數情形 channel 是 4的倍數 （否則將使用 NCHW 佈局下的預設實現）。同時，為了更好的記憶體區域性性，我們希望沿著最內層的軸進行累加。考慮到這些因素，我們使用一種自定義的佈局。

在 CUDA int8 conv2d 中，更加經驗我們選用了 NCHW4c 作為資料佈局，OIHW4o4i 作為卷積核的佈局。實際上，我們的實現能夠很容易地推廣到 NCHW[x]c 和 OIHW[x]o[x]i 這種更一般的情況。其中，x 是任意可被 4 整除的正整數。在我們選用的 NCHW4c 資料佈局中，我們沿 channel 維度將每 4 個元素打包放到最內層。類似地，我們將卷積核中，我們將對應輸入和輸出 channel 的維度也分別進行打包並移到最內層。這種卷積核佈局使得輸出具有和輸入相同的佈局。因此，我們不需要在每層之間插入不必要的佈局轉換。

圖2 展示了計算 2d卷積輸出的一個元素的過程。在我們選用的佈局(NCHW4c和OIHW4o4i）中，按 NCHW 或 OIHW 四個維度索引得到的結果是打包後的元素。圖2中間，卷積核部分的每一列來自於原來卷積核中不同的filter。我們用 dp4a 計算中間這幅圖packed kernel部分每一行與 packed input 的點積並累加到輸出結果上。

在指定卷積層的佈局之後，在 Relay 的 AlterOpLayout pass中，我們會自動地將其他的運算子如張量加法、啟用函式等轉換到所選的佈局。對卷積核的佈局轉換可以離線地提前計算。我們可以讓整個網路在新的資料佈局下執行，而沒有額外的開銷。

設計自動優化的搜尋空間

我們所實現的量化版本的運算子達到好的效能的關鍵是使用基於機器學習的自動優化。在這裡，我們需要考慮一個問題：如何設計一個有效地 schedule 的搜尋空間？一個有效的 schedule 搜尋空間是指自動調優過程能夠在合理的的迭代次數下達到理想的效能。通常來說，我們需要儘量地讓搜尋所使用的模版更加靈活，儘可能地覆蓋搜尋空間中不同的配置。另一方面，我們也需要藉助效能優化上的先驗知識。例如，把資料快取到 shared memory 是 CUDA 程式設計的一個常用技巧。因此，我們在使用 shared memory 的同時，讓機器學習演算法來自動的選擇 shared memory 大小的最優分配。我們也手動地做了一些 tiling 例如把一些維度按 4 或 16的因子進行分割來以利用向量化的記憶體訪問。

在 int8 2d 卷積演算法中，我們設計的搜尋空間包含了一組可調的選項，例如 tile size，融合多層迴圈時的選擇，loop unrolling （迴圈展開） 及 double buffering （雙快取）的設定。TVM 中，CUDA平臺量化版本的 conv2d 和 dense 註冊在 int8 這一 template key下。在做自動調優的時候，我們可以在建立任務的時候設定 template key 來選擇這一實現。關於自動調優的細節可以檢視這篇AutoTVM教程.

總體工作流程

TVM 提供了簡單的工作流程。只需幾步，即可量化從其他框架訓練好的模型、自動優化模型中的運算子（基於AutoTVM），並部署到不同的裝置上。

首先，我們使用 Relay 前端匯入現有的模型。在此我們使用一個輸入形狀為 (1, 3, 224, 224) 的MXNet模型為例。

sym, arg_params, aux_params = mxnet.model.load_checkpoint(model_path, epoch)
net, params = relay.from_mxnet(sym, shape={'data': (1, 3, 224, 224)}, arg_params=arg_params, aux_params=aux_params)

接下來，我們直接呼叫 Relay 的量化 API 生成量化後的模型。

net = relay.quantize.quantize(net, params=params)

得到量化模型之後，我們使用 AutoTVM 對模型中使用的運算子建立調優任務並進行自動優化。該部分可以參考AutoTVM教程。

最後，我們編譯最終得到的模型，即可在量化模式下執行。

with relay.build_config(opt_level=3):
graph, lib, params = relay.build(net, target)

relay.build 的返回值可以直接用於部署。我們可以直接在本地GPU上執行，或 通過 RPC 部署到遠端裝置 .

效能測試

為了驗證 TVM 中量化版本的運算子的效能，我們在一些常用的模型如 VGG-19、ResNet-50、Inception V3 上進行測試。此外，我們還測試了DRN-C-26, ResNeXt-50, 和 Deformable ConvNets 中的 DCN-ResNet-101 來比較新出現的模型上的效能。這些模型包含一些較少使用的運算，如 dilated 卷積、 group 卷積、deformable 卷積等。我們選用 NVIDIA TensorRT 作為 baseline. 提供的 MXNet 1.4 + cuDNN 7.3 在 float32 模式下的資料可以作為量化加速的參考。該測試是在一張 NVIDIA GTX 1080 上進行的。我們報告在 batch size = 1 和 16下，平均每張輸入圖片的執行時間。

如圖1所示，相比於 float32 的模型，使用 TVM 量化後能夠達到最高8倍的加速。在VGG、ResNet等標準的 CNN 模型上，TVM 達到了近似於 TensorRT 的 state-of-the-art 結果。

TVM 在新出現的模型上也達到了很好的結果。在 ResNeXt 和 DCN-ResNet-101 上 TVM 有顯著的效能提升。因為沒有 deformable 卷積的官方實現，我們暫無 TensorRT 在 DCN-ResNet-101 上的結果。可以看出，自動優化使得 TVM 能夠簡單且靈活地支援和優化新出現的模型。