模型壓縮明珠:二值化網路

語言: CN / TW / HK

二值化網路(BNN)

  • 老闆:量化到INT8又怎麼樣!還不夠小!我要把AI模型放在耳機手錶裡面!!
  • 員工:那我們用二值化網路!!一切都是0和1!!

二值化網路跟低位元量化一樣,目的是讓模型更小,小到有著最為極端的壓縮率和極低的計算量。那什麼是二值呢?二值指的是僅僅使用+1和-1(或者是0和1)兩個值,來表示權重和啟用的神經網路。

相比於全精度(FP32)表示的神經網路,二值化可以用XNOR(邏輯電路中的異或非門)或者是簡單的計數操作(pop Count),極其簡單的組合來代替FP32的乘和累加等複雜的運算來實現卷積操作,從而節省了大量的記憶體和計算,大大方便了模型在資源受限裝置上的部署。

BNN由於可以實現極高的壓縮比和加速效果,所以是推動以深度神經網路為代表的人工智慧模型,在資源受限和功耗受限的移動裝置上落地應用的一門非常具有潛力的技術啦。

不過呢,目前BNN仍然存在著很多不足,例如模型精度仍然比全精度低,無法有效地泛化到更復雜的任務上,依賴於特定的硬體架構和軟體框架等。但同時也能看到,BNN從2016年首次提出時,ImageNet上只有27%的Top1準確率,到2020年提出的ReActNet-C有著71.4%準確率的進步!

下面來看看BNN在AI系統全棧/AI框架中的一個位置,橙色標籤的地方是BNN所在的位置。從圖中可以知道,在表達層API需要提供二值化網路模型使用到的API介面;接著在中間邊一層和runtime其實沒有太多的區別,最主要是底層需要提供二值化相關的運算元或者是二值化推理的專用硬體電路。

1. BNN基本介紹

BNN最初由Yoshua Bengio[1]在2016年的時候首次提出來,論文中表示使用隨機梯度下降的方式來訓練帶有二值化的權重weight和啟用act引數的神經網路模型。

1.1 前向計算

為了解決二值化權重weight計算中梯度的傳遞問題,論文中提出在訓練過程中保持一個實值(FP32)的權重,然後使用一個sign函式來獲得二值化的權重引數。 [公式] 為FP32, [公式] 為二值化後的值:

[公式]

其中sign函式為只要輸入的引數大於等於0則為1,否則都是-1:

[公式]

下面圖中上面的是二值化權重和二值化輸入的一個3X3卷積操作,二值化操作則是對卷積核和輸入資料的視窗進行平鋪展開,然後進行XNOR操作,接著進行bit count計數得到卷積結果。

1.2 反向傳播

跟感知量化訓練的方式類似,sign函式在0處不可導,導數為0時會遇到沒有辦法計算梯度,於是論文中提出了直通估計 (straight through estimator STE),即當梯度傳遞遇到sign函式時,直接跳過這個函式:

[公式]

使用了直通估計STE之後,可以使用與全精度神經網路相同的梯度下降方法直接訓練二值神經網路。權重引數可以使用常見的優化器來進行更新,同時考慮到訓練過程權重引數是沒有進行截斷的,這樣權重引數可能會一直累加到特別大的值,從而與二值化的權重之間的量化誤差越來越大,因此論文對訓練過程中的權重增加截斷函式,將其限制在-1和+1之間,這樣使得訓練過程中權重引數和二值化的權重引數的誤差偏差不會過大:

[公式]

既然對FP32訓練的過程有嵌入式的修改,那麼肯定會導致訓練的時間更長了,而且最終的實驗結果準確率沒有FP32那麼高,有什麼用嗎?

這裡翻譯翻譯,其實最大的作用是前向的圖所示,可以1 bit的異或非和pop count操作,來代替FP32的卷積進行乘和累加操作,在實際模型部署和推理的時候,不僅能減少32倍的記憶體引數儲存,還能跑的比馬都快!

3. BNN網路結構

近年來,各種各樣的二值神經網路方法被提出,從使用預定義函式直接對權重和輸入進行量化的樸素二值化方法,到使用基於多種角度和技術的基於優化的二值化方法,其中包括通過最小化量化誤差來近似全精度值、通過修改網路損失函式來限制權重、和通過減小梯度誤差來學習離散引數。

其中北京航空航天大學最新綜述文章 Binary Neural Networks: A Survey 已經對很多二值化的網路模型寫了個比較好的綜述了,下面ZOMI簡單介紹個我覺得比較精彩的BNN網路模型結構。

4. 硬體實現

從二值化網路的流程來看,BNN的主要加速原因就是用XNOR與Pop Count操作來代替了傳統卷積演算法中,使用昂貴的乘法-累加MAC操作。

而通用的x86計算架構,基本上都是對FP32全精度型別資料的計算,進行底層的硬體和編譯執行優化,所以直接將BNN部署在通用的x86計算平臺收益其實並不明顯,甚至可能不僅沒有加速效果,甚至比同等的FP32網路模型執行還慢。

下面分別看在ARM CPU和FPGA平臺的一個簡單分析。

ARM CPU

BNN其實目前聚焦部署在移動端。BMXNet 2017[3]由來自德國Hasso Plattner Institute的研究員Haojin Yang等開發,其為一個基於MXNet的二值化開源框架。支援使用CuDNN進行訓練,並使用二值運算子XNOR和Pop Count做推理。不足之處是二值化核心並未經過專門調優,因此在ARM CPU上的速度表現並不突出。

Dabnn 2019[4]由京東AI研究院推出,基於彙編調優後的BNN推理工具。它有效提升了BNN在ARM框架端上的推理速度,但是這個工具並不能用於模型訓練,需要使用其它工具進行訓練。

 

BMXNet-v2 2019[9],Bethge和Yang等開源了支援Gluon API的第二版。該框架採用的一系列改進大大降低了模型訓練難度,大大減小了MXNet同步的成本。第二版不僅提升了效率,同時繼續支援模型的壓縮和二值化推理,可將模型部署在多種邊緣裝置上。

FPGA和ASIC

相比於傳統的CPU,FPGA在硬體架構設計方面很靈活,可以支援bits-wise的高效運算,而且功耗很低,其對應的終極版的專用積體電路ASIC更是如此,可以比FPGA運算更高效,更節能。

目前用FPGA設計AI加速器基本都是以Xilinx的器件和開發工具為主,而他們也為二值化神經網路專門設計了一款架構FINN,開發者可以利用高層次綜合工具(HLS),用C語言進行開發,直接將二值化模型部署到FPGA上。

總結

雖然BNN在近5年來已經取得了很大的進步啦,但是比較大的問題就是精度的損失仍然是個頭痛的問題,尤其是對於大型網路和資料集而言。主要原因可能包括:

1)目前還沒有SOTA的二值化網路模型,不確定什麼樣的網路結構才合適做二值化;

2)即使有用於二值化的梯度估計或近似函式,在離散空間中優化二值網路是一個難題。

另外,隨著移動裝置廣泛使用,將出現更多針對這些應用進行的研究工作,以實現不同的任務以及模型在不同硬體上的部署。例如在耳機進行觸感分析、點選分析這種大部分都是對訊號進行分類操作,其實不需要大模型,反而這個時候二值化網路,能夠很好地對訊號資料進行高精度分類,去判斷到底是雙擊停止播放音樂還是拖動放大聲音。

最後,對可解釋機器學習的研究表明,神經網路的推理中存在關鍵路徑,並且不同的網路結構遵循不同的模式。因此,根據層的重要性設計混合精度策略,並設計出對二值神經網路的資訊流友好的新網路結構,也具有重要意義。

參考文獻

  • [1] Courbariaux, Matthieu, et al. "Binarized neural networks: Training deep neural networks with weights and activations constrained to+ 1 or-1." arXiv preprint arXiv:1602.02830 (2016).
  • [2] Qin, Haotong, et al. "Binary neural networks: A survey." Pattern Recognition 105 (2020): 107281.
  • [3] Yang, Haojin, et al. "Bmxnet: An open-source binary neural network implementation based on mxnet." Proceedings of the 25th ACM international conference on Multimedia. 2017.
  • [4] Zhang, Jianhao, et al. "dabnn: A super fast inference framework for binary neural networks on arm devices." Proceedings of the 27th ACM international conference on multimedia. 2019.
  • [5] https://zhuanlan.zhihu.com/p/27
  •