來源:新浪證券
背景
近年來隨著邊緣計算和物聯網的興起與發展,許多移動終端(比如手機)成為了深度學習應用的承載平臺,甚至出現了各式各樣專用的神經網絡計算晶片。由於這些設備往往對計算資源和能耗有較大限制,因此在高性能伺服器上訓練得到的神經網絡模型需要進行裁剪以縮小內存佔用、提升計算速度後,才能較好地在這些平臺上運行。
一種最直觀的裁剪方式就是用更少位數的數值類型來存儲網絡參數,比如常見的做法是將 32 位浮點數模型轉換成 8 位整數模型,模型大小減少為 1/4,而運行在特定的設備上其計算速度也能提升為 2~4 倍,這種模型轉換方式叫做量化(Quantization)。
量化的目的是為了追求極致的推理計算速度,為此捨棄了數值表示的精度,直覺上會帶來較大的模型掉點,但是在使用一系列精細的量化處理之後,其在推理時的掉點可以變得微乎其微,並能支持正常的部署應用。
原理
實現量化的算法多種多樣,一般按照代價從低到高可以分為以下四種:
Type1 和 Type2 由於是在模型浮點模型訓練之後介入,無需大量訓練數據,故而轉換代價更低,被稱為後量化(Post Quantization),區別在於是否需要小批量數據來校準(Calibration);
Type3 和 Type4 則需要在浮點模型訓練時就插入一些假量化(FakeQuantize)算子,模擬量化過程中數值截斷後精度降低的情形,故而稱為量化感知訓練(Quantization Aware Training, QAT)。
以常用的 Type3 為例,一個完整的量化流程分為三階段:(1)以一個訓練完畢的浮點模型(稱為 Float 模型)為起點;(2)包含假量化算子的用浮點操作來模擬量化過程的新模型(Quantized-Float 模型或 QFloat 模型);(3)可以直接在終端設備上運行的模型(Quantized 模型,簡稱 Q 模型)。
由於三者的精度一般是 Float > QFloat > Q ,故量化算法也就分為兩步:
· 拉近 QFloat 和 Q:這樣訓練階段的精度可以作為最終 Q 精度的代理指標,這一階段偏工程;
· 拔高 QFloat 逼近 Float:這樣就可以將量化模型性能儘可能恢復到 Float 的精度,這一階段偏算法。
第一步在MegEngine框架的「訓練推理一體化」特性下得到了保證,而第二步則取決於不同的量化算法。
儘管不同量化算法可能在假量化的具體實現上有所區別,但是一般都會有一個「截斷」的操作,即把數值範圍較大的浮點數轉換成數值範圍較小的整數類型,比如下圖,輸入一個[-1, 1)範圍的浮點數,如果轉換為 4 位整型,則最多只能表示 2^4 個值,所以需要將輸入的範圍劃分為16段,每段對應一個固定的輸出值,這樣就形成了一個類似分段函數的圖像,計算公式為:
[公式]
另外,由於分段函數在分段點沒有梯度,所以為了使假量化操作不影響梯度回傳,就需要模擬一個梯度,最簡單的方法就是用y=x來模擬這一分段函數,事實證明這麼做也是有效的,這種經典的操作被稱為「Straight-Through-Estimator」(STE)。
工程
量化部分作為模型推理部署的重要步驟,是業界在大規模工業應用當中極為關注的部分,它在 MegEngine 的底層優化中佔了很大比重。在目前開源的版本裡,針對三大平臺(X86、CUDA、ARM),MegEngine都有非常詳細的支持,尤其是ARM平臺。
一般在通用計算平臺上,浮點計算是最常用的計算方式,所以大部分指令也是針對浮點計算的,這使得量化模型所需的定點計算性能往往並不理想,這就需要針對各個平臺優化其定點計算的性能。
ARM 平臺
ARM平臺一般是指手機移動端,其系統架構和底層指令都不同於我們熟知的電腦CPU,而隨著架構的變遷,不同架構之間的指令也存在不兼容的問題。為此,MegEngine針對ARM v8.2前後版本分別實現了不同的優化:
· ARM v8.2 主要的特性是提供了新的引入了新的 fp16 運算和 int8 dot 指令,MegEngine基於此進行一系列細節優化(細節:四個int8放到一個128寄存器的32分塊裡一起算),最終實現了比浮點版本快2~3倍的速度提升
· 而對於v8.2之前的ARM處理器,MegEngine則通過對Conv使用nchw44的layout和細粒度優化,並創新性地使用了int8(而非傳統的int6)下的winograd算法來加速Conv計算,最使實現能夠和浮點運算媲美的速度。
CUDA 平臺
CUDA 平臺是指 NVIDIA 旗下 GPU 平臺,由於提供 CUDNN 和 Toolkit 系列接口以及 TensorRT 專用推理庫,大部分算子可以使用官方優化,而 MegEngine 則在此基礎上進行了更多細節的優化,比如如何更好地利用 GPU 的TensorCore 進行加速,不同型號之間一些差異的處理等,最終效果根據不同模型也有非常明顯的推理加速。
X86 平臺
X86 平臺是指 Intel CPU 平臺,近年來隨著深度學習的發展,其也慢慢提供了針對定點運算更多的支持。
· 在新一代至強(Xeon)處理器上,通過使用 VNNI(Vector Neural Network Instructions)指令,MegEngine 將 CPU 的 int8 推理性能優化到了浮點性能的 2~3 倍。
· 而對於不支持 VNNI 指令的 CPU,一般只提供最低 int16 的數值類型支持,則通過使用 AVX2(Advanced Vector Extensions)這一向量格式,實現了 int8 推理性能與浮點性能持平。
以上是對各個平臺推理加速效果的整體介紹,更多更細節的介紹可以期待之後的系列文章。
使用
除了底層實現上的加速與優化,在 Python 側訓練部分,MegEngine對接口也有很多細節設計,使得整體代碼邏輯清晰簡潔。
我們在 Module 中額外引入了兩個基類:QATModule、QuantizedModule 。分別代表上文提及的帶假量化算子的 QFloat 模型與 Q 模型,並提供普通 Module → QATModule → QuantizedModule 三階段的轉換接口。各個版本的算子是一一對應的,且通過合理的類繼承免除了大量算子實現中的冗餘代碼,清晰簡潔。
如上圖,用戶首先在普通 Module 上進行正常的模型訓練工作。訓練結束後可以轉換至 QFloat 模型上,通過配置不同的 Observer 和假量化算子來選擇不同的量化參數 scale 獲取方式,從而選擇進行 QAT 或 Calibration 後量化。之後可以再轉換至 Q 模型上,通過 trace.dump 接口就可以直接導出進行部署。
針對推理優化中常用的算子融合,MegEngine 提供了一系列已 fuse 好的 Module,其對應的 QuantizedModule 版本都會直接調用底層實現好的融合算子(比如 conv_bias)。
這樣實現的缺點在於用戶在使用時需要修改原先的網絡結構,使用 fuse 好的 Module 搭建網絡,而好處則是用戶能更直接地控制網絡如何轉換,比如同時存在需要 fuse 和不需要 fuse 的 Conv 算子,相比提供一個冗長的白名單,我們更傾向於在網絡結構中顯式地控制,而一些默認會進行轉換的算子,也可以通過 disable_quantize 方法來控制其不進行轉換。
另外我們還明確了假量化算子(FakeQuantize)和Observer的職責,前者將主要負責對輸入進行截斷處理的計算部分,而後者則只會記錄輸入的值,不會改變輸出,符合 Observer 的語義。
在配置使用上,用戶需要顯式指定針對 weight、activation 分別使用哪種 Observer 和 FakeQuantize,比如:
這樣的好處在於,用戶可以控制每一處量化過程的細節,可以分別採用不同量化算子和數值類型。
下文簡單說明一下在 MegEngine 中轉換一個 ResNet 網絡的全流程代碼:
Float → QFloat:
QFloat → Q 並導出用於部署:
本文簡單介紹了神經網絡模型實際應用在移動平臺必不可少的一步——量化,以及天元(MegEngine )在量化上做的一些工作:包括底層針對不同平臺的一些優化效果,在用戶接口使用上的一些設計理念。
天元(MegEngine)相信,通過簡潔清晰的接口設計與極致的性能優化,「深度學習,簡單開發」將不僅惠及曠視自身,也能便利所有的研究者,開發者。
參考文獻
[1] Moons, B., Goetschalckx, K., Van Berckelaer, N., & Verhelst, M. (2017, October). Minimum energy quantized neural networks. In 2017 51st Asilomar Conference on Signals, Systems, and Computers (pp. 1921-1925). IEEE.
[2] Jacob, B., Kligys, S., Chen, B., Zhu, M., Tang, M., Howard, A., ... & Kalenichenko, D. (2018). Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (pp. 2704-2713).
[3] Zhou, A., Yao, A., Guo, Y., Xu, L., & Chen, Y. (2017). Incremental network quantization: Towards lossless cnns with low-precision weights. arXiv preprint arXiv:1702.03044.
[4] Li, F., Zhang, B., & Liu, B. (2016). Ternary weight networks. arXiv preprint arXiv:1605.04711.
[5] Rastegari, M., Ordonez, V., Redmon, J., & Farhadi, A. (2016, October). Xnor-net: Imagenet classification using binary convolutional neural networks. In European conference on computer vision (pp. 525-542). Springer, Cham.