8比特數值也能訓練模型?商湯提出訓練加速新算法
導讀:在CVPR 2020上,商湯研究院鏈接與編譯團隊、高性能計算團隊和北航劉祥龍老師團隊合作提出了用于加速卷積神經網絡訓練過程的INT8訓練技術。該工作通過將網絡的輸入、權重和梯度量化到8比特來加速網絡的前向傳播和反向傳播過程,縮短卷積神經網絡訓練時間。論文觀察到梯度的獨特分布給量化訓練帶來了極大挑戰,為了解決梯度量化帶來的精度損失和不穩定問題,該論文進行了量化訓練收斂穩定性的理論分析并基于此提出了誤差敏感的學習率調節和基于方向自適應的梯度截斷方法。同時為了保證更高的加速比,該論文還提出使用周期更新、量化卷積融合等技術來減少量化操作帶來的時間開銷。應用了上述方法之后,INT8訓練在圖像分類任務和檢測任務上都僅僅損失微小的精度,且訓練過程相比浮點訓練加速了22%。
動機與背景
卷積神經網絡被廣泛應用在多種計算機視覺任務中并且取得了優異的精度。由于擁有龐大的參數量,訓練和部署卷積神經網絡需要耗費大量計算資源和漫長的訓練時間,如何用更少資源訓練卷積神經網絡一直是一個學術研究熱點,也是工業界關心的話題。
神經網絡量化技術是一種使用定點計算代替浮點的加速技術,目前被廣泛地應用在神經網絡部署中,可以極大地提升部署速度,并降低內存資源占用。現有很多工作均表明將網絡前向過程的浮點計算替換成INT8計算,不會帶來明顯的精度下降[1][2]。
下圖展示了現代神經網絡加速芯片對于不同精度計算的理論計算峰值對比,可以看到,INT8算力相比于FP32和FP/INT16均能有超過2倍峰值性能提升。
當考慮將神經網絡量化技術應用在卷積神經網絡訓練中時,為了加速卷積的反向梯度傳播過程,不得不對梯度進行量化操作。在將浮點的梯度量化到INT8數值范圍內之后,訓練過程變得極其不穩定,并且收斂到非常差的精度。如何解決量化梯度給訓練帶來的收斂穩定性問題,是十分重要的問題。與此同時,在提升訓練精度的同時,也不應當進入過多額外的計算,否則加速效果將會大打折扣。
一方面是高效的計算峰值保障,一方面是困難重重的算法設計,這是INT8訓練技術的機遇與挑戰。
何為INT8訓練
標準的線性量化操作指的是,將一個浮點張量(tensor)進行線性映射,變換到整數空間中[3]。這個整數空間的大小由于量化比特數來決定,比如常見的8bit量化數,就有256個取值,本文中使用的是對稱量化,因此量化數的取值是從-128到127。具體公式如下,其中x是被量化的數據,q是量化后的數據,s是量化系數,clip是截斷函數:
在8bit的場景里,截斷函數和量化系數的計算公式如下:
為了降低量化帶來的誤差,一個常見做法是對取整過程進行隨機化,使得取整函數從期望上更接近原始的數,具體隨機取整的公式如下:
相反的,將8bit量化數變換回浮點的過程稱之為反量化。反量化公式如下所示,其中q為量化計算結果,s為量化系數,為反量化后的結果。
上圖的上半部分展示了標準的卷積神經網絡量化計算前向過程,該過程被廣泛應用在INT8部署加速中。在卷積計算之前,量化器會對輸入和權重進行量化操作,將浮點數量化到8bit數值上,通過INT8卷積計算核心,即可完成一次INT8前向計算,最終將求和得到的32bit數進行反量化操作回算到浮點數域中,以供給下一層計算使用。
INT8訓練的一個核心的加速點在于卷積計算的反向過程,上圖展示了INT8訓練中卷積計算在反向傳播過程中的計算細節。在卷積的反向梯度傳播過程,同樣的將梯度進行浮點量化操作,不過為了降低量化的誤差,針對梯度的量化采用了隨機取整操作。通過INT8的反向卷積計算核心,可以得到下一層所需的回傳梯度,以及當前層的權重所需的梯度。由于INT8反向卷積輸出的是32bit數,與前傳類似,需要引入一次反量化操作,將32bit數反算回到浮點數域中。
梯度為何難以量化
為什么對梯度進行量化會給網絡訓練帶來如此大的影響?我們可以觀察訓練過程中的梯度分布情況來進一步的分析。
通過圖(a)中對比梯度和輸入、權重的分布,可以發現:梯度分布相比輸入和權重分布更加尖銳,同時范圍更大。相比于輸入和權重,梯度有更多的值集中在0附近,但同時梯度還有許多較大值,讓梯度的分布范圍變得相當廣,這些特征都會導致梯度量化的量化誤差比輸入和權重更大。
圖(b)展示的是layers16隨著訓練,其梯度從epoch 0到epoch 300的變化情況。從中可以看出,隨著訓練的進行,梯度分布越變得更加尖銳,同時仍然保持著較廣的分布范圍,這意味著梯度量化的誤差會隨著訓練的進行變得越來越大。
梯度的分布隨網絡深度變化情況從圖(c)中可以看出。很容易發現,卷積層的深度越淺,梯度分布越尖銳,這也會導致梯度量化的誤差更大。
從圖(d)中可以看出卷積的結構也會影響梯度分布,對于MobileNetV2來說,conv2為depthwise卷積其相比conv1和conv3具有更加尖銳的分布。
由于卷積神經網絡的梯度具有如上四個特點,所以當我們直接在訓練中對梯度進行量化時,訓練精度非常容易出現突發的崩潰情況。下圖展示了在CIFAR-10數據集上進行實驗的精度和損失函數變化曲線,以MobileNetv2在CIFAR-10數據集上訓練為例,其訓練的精度曲線和loss曲線如下圖,從圖中可以發現INT8訓練的loss在訓練初期正常下降,但隨后迅速上升,對應的精度也不斷下降。
是什么影響了收斂穩定性
根據以上的觀察和初步啟發,我們希望通過理論的分析和推導,對量化訓練的收斂穩定性進行建模。根據Adam等相關論文的經驗和優化理論中的Regret analysis,不失一般性地定義R(T)為
其中f是損失函數,t是訓練輪數,T是訓練總輪數,
其中f是損失函數,t是訓練輪數,T是訓練總輪數,為t輪的權重,
是最優權重。
基于以下兩個樸素的假設:
通過推導證明可以得到:
其中為t輪的學習率,
d為權重的維度,
為t輪的量化誤差,
是t輪的量化后梯度。
為了確保網絡能夠穩定收斂,在T變大時需要能夠達到足夠小。通過上式可以發現,在T趨于無窮大時,第(1)項可以忽略不計,主要考慮減小第(2)項和第(3)項。
我們發現,第(2)項與量化誤差正相關,第(3)項與學習率以及量化后的梯度大小有關。
因此我們不難得到兩個直觀的提升訓練收斂穩定性的策略:
通過調節量化函數中的截斷減小量化誤差
通過適當調低學習率來提高量化訓練精度
主要方法
依據以上分析,我們針對量化誤差和學習率提出了基于方向自適應的梯度截斷和誤差敏感的學習率調節兩個方法來解決量化訓練帶來的精度損失問題。同時,為了減少量化操作帶來的額外開銷,本文還提出了周期更新和量化卷積融合的方法。
1、 基于方向自適應的梯度截斷:調整截斷值,讓梯度方向保持正確。
為了最小化量化誤差,之前有很多研究提出優化截斷值的方法,其中就有研究提出通過假設數據分布直接求解最優截斷值。但是已有的研究都針對于權重量化的截斷值進行優化。就如本文觀察所顯示,梯度的分布特征與權重區別較大,無法直接使用。本文通過KS檢驗發現梯度的分布并不符合常見的高斯分布、拉普拉斯分布和學生t分布,因此很難通過假設梯度分布來直接求解最優的截斷值。
基于以上的分析,本文采用梯度下降的方法來自適應地學習最優截斷值,常見的目標函數有均方誤差函數,但是由于梯度的分布特征,均方誤差的大小會受到梯度的影響,影響優化過程;同時對于梯度來說,均方誤差并不能很好地體現梯度的量化誤差對于優化過程的影響,因此本文提出使用能夠體現梯度方向的余弦距離來衡量梯度的量化誤差,并以余弦距離為目標函數來優化求解最優截斷值。余弦距離
定義如下:
其中,是梯度,
是量化后的梯度。
2、 誤差敏感的學習率調節:在錯誤的方向上盡量少更新。
根據上述的理論分析,降低學習率能夠有助于模型量化訓練的收斂。針對學習率的調整,本文提出誤差敏感的學習率調節方法,使用學習率系數對原學習率進行調整,學習率系數與余弦距離負相關,學習率系數
定義如下:
其中α和β是超參數,用于控制衰減程度和調節下界。
3、 周期更新:降低由于統計而帶來的額外計算耗時
由于量化操作需要的統計數據范圍和計算截斷值等操作十分耗時,為了減少這些操作的時間開銷,本文采用周期更新的方式,周期性地統計數據范圍和計算截斷值。通過周期更新的方法能夠有效地提高減少因量化引入的額外時間開銷。下表為ResNet50在ImageNet數據集上不同周期的單次訓練時間統計表。
4、 量化卷積融合:減少訪存次數、節省cuda kernel launch次數
通過將量化和反量化操作融合入卷積計算的CUDA核函數里,可以減少一次數據的訪存,有效地減少量化和反量化操作的時間開銷。
實驗結果
圖像分類任務:本文在CIFAR10和ImageNet等圖像分類數據集進行INT8訓練實驗。從下表結果中可以看出,在大多數網絡結構中均取得了比現有最好方法更優的精度,并且首次在MobileNet、Inception等網絡上進行量化訓練實驗,精度損失也在1.5%以內。
目標檢測任務:同時,本文也首次嘗試在PASCAL和COCO等目標檢測數據集上進行INT8訓練實驗,精度損失也在2%以內。
已有的少量探究梯度量化的論文[4]均未報告算法在實際訓練任務中的真實加速性能,為了最大限度將方法實用化,本文在 GeForce GTX1080TI顯卡上編寫并優化了用于支持INT8訓練的卷積前向和后向計算核心。實測結果表明,使用INT8卷積計算的前向和后向過程相比于浮點計算有明顯的加速,其中前向過程平均加速1.63倍,后向過程平均加速1.94倍。如下圖所示:
同時,本文在實際訓練過程中進行了完整的端到端測試,可以看到,INT8訓練可以將ResNet50的一輪訓練過程從0.360秒降低到0.293秒,整體訓練過程提速了22%。
招聘
商湯研究院-鏈接與編譯團隊致力于通過System+AI技術打造頂尖的深度學習核心引擎。開發的模型訓練和模型部署工具鏈已服務于公司多個核心業務。團隊在量化模型的在線/離線生產、部署對齊、加速訓練等方向有著明確的技術規劃。
團隊目前正在招聘模型量化方向的系統算法研究員/實習生,歡迎感興趣的小伙伴投遞簡歷到yufengwei@sensetime.com。
傳送門
論文地址:https://arxiv.org/pdf/1912.12607.pdf
論文講解:近期將有線上直播的論文講解及專題交流課程,歡迎點擊下方“商湯泰坦公開課”課程專題鏈接進行報名
報名鏈接:https://jinshuju.net/f/RKOLf9