【論文閱讀】Mixed Precision Traning

【GiantPandaCV導語】混合精度是一個非常簡單并且實用的技術，由百度和谷歌聯(lián)合發(fā)表于ICLR2018，可以讓模型以半精度的方式訓練模型，既能夠降低顯存占用，又可以保持精度。這篇文章不是最先提出使用更低精度來進行訓練，但是其影響力比較深遠，很多現(xiàn)在的方案都是基于這篇文章設計的。

1. 摘要

提高網(wǎng)絡模型的大小可以有效提升準確了，但是也增加了內存的壓力和計算力的需求。本論文引入了半精度來訓練深度神經(jīng)網(wǎng)絡，在不損失精度、不需要更改超參數(shù)的情況下，幾乎可以減半內存占用。

權重、激活、梯度都是以IEEE半精度的方式存儲的，由于半精度的表示范圍要比單精度要小，本文提出了三種方法來防止關鍵信息的丟失。

• 推薦維護一個單精度版本的權重weights，用于累積梯度。
• 提出使用Loss scale來保證小的梯度值不會下溢出。
• 在加載到內存之前，就將模型轉化到半精度來進行加速。

2. 介紹

大的模型往往需要更多的計算量、更多的內存資源來訓練，如果想要降低使用的資源，可以通過降低精度的方法。通常情況下，一個程序運行速度的快慢主要取決于三個因素：

• 算法帶寬
• 內存帶寬
• 時延

使用半精度減少了一半bit，可以影響算法帶寬和內存帶寬，好處有：

• 由于是用更少的bit代替，占用空間會減小差不多一半
• 對于處理器來說也是一個好消息，使用半精度可以提高處理器吞吐量；在當時的GPU上，半精度的吞吐量可以達到單精度的2-8倍。

現(xiàn)在大部分深度學習訓練系統(tǒng)使用單精度FP32的格式進行存儲和訓練。本文采用了IEEE的半精度格式(FP16)。但是由于FP16可以表示的范圍要比FP32更為狹窄，所以引入了一些技術來規(guī)避模型的精度損失。

• 保留一份FP32的主備份。
• 使用loss scale來避免梯度過小。
• FP16計算但是用FP32進行累加。

3. 實現(xiàn)

下面詳細講解第二節(jié)提到的三種技術。

3.1 FP32的主備份

在混合精度訓練中，權重、激活、梯度都采用的是半精度存儲的。為了匹配單精度網(wǎng)絡的精度，在優(yōu)化的整個過程中，需要拷貝一份單精度模型FP32作為主備份，而訓練過程中使用的是FP16進行計算。

這樣做的原因有兩個：

第一：溢出錯誤：更新的時候是學習率乘以梯度，這個值可能非常小，超過半精度最小范圍()，就會導致下溢出。如下圖展示了梯度和FP16的表示范圍，大約有5%的梯度已經(jīng)下溢出，如果直接對半精度表示的模型更新，這樣模型的精度肯定有所損失；而對單精度的模型進行更新，就就可以規(guī)避下溢出的問題。

• 第二：舍入誤差：指的是當梯度過小，小于當前區(qū)間內的最小間隔，那么梯度更新可能會失敗。想要詳細了解可以自行查看wiki上的詳解。同樣的，使用FP32單精度可以盡可能減少舍入誤差的影響。

3.2 Loss Scale

下圖展示的是激活值的范圍和FP16表示范圍之間的差異，50%以上的值都集中在不可表示的范圍中，而右側FP16可表示范圍卻空蕩蕩的，那么一個簡單的想法就是向右平移，將激活平移到FP16可表示范圍內。比如說可以將激活乘以8，這樣可以平移3位，這樣表示最低范圍就從到了, 因為激活值低于的部分對模型訓練來說不重要，所以這樣操作就可以達到FP32的精度。

那實際上是如何實現(xiàn)平移呢？很簡單，就是在loss基礎上乘以一個很大的系數(shù)loss scale, 那么由于鏈式法則，可以確保梯度更新的時候也采用相同的尺度進行縮放，最后在更新梯度之前再除以loss scale， 對FP32的單精度備份進行更新，就完成了。

還有一個問題，就是loss scale值的選取，最簡單的方法就是選擇一個固定值作為loss scale,可以訓練一系列網(wǎng)絡縮放因子的設置可以從8到32k。固定的放縮因子選取是需要一定經(jīng)驗的，如果梯度的數(shù)據(jù)可以獲取，那就選取一個值可以讓其梯度的值小于65504（FP16可以表示的最大值）。如果選取loss scale值過大，可能會上溢出，這種情況可以通過檢查梯度來判斷，如果發(fā)生溢出，跳過權重更新，進行下一次迭代。

3.3 運算精度

在神經(jīng)網(wǎng)絡中的運算可以劃分為三類：向量點乘（dot-products）、reduction和逐點操作。這三種操作在遇到半精度的時候，會有不同的處理方法。

為了保持模型精度，需要將向量點乘的部分乘積累加為FP32，然后再轉換為FP16。如果直接用FP16會導致精度損失，不能達到和原來一樣的精度?，F(xiàn)在很多支持Tensor Cores的GPU中已經(jīng)實現(xiàn)了這樣的操作。

Reduction操作（比如求很多向量元素的和）需要先轉為FP32, 這樣的操作大部分出現(xiàn)在batch norm或者softmax層中。這兩個層都是從內存中讀寫FP16的向量，執(zhí)行運算的時候要使用FP32格式。

逐點操作，如非線性和元素間矩陣乘法，內存占用有限，運算精度不影響這些操作的速度。

4. 結果

在分類任務（ILSVRC/ImageNet）、檢測任務（VOC2007）、語音識別（English-Mandarin）、機器翻譯（English-Frech）、語言模型（1 billion word）、DCGAN這幾個任務進行了實驗，這篇解讀僅展示分類任務和檢測任務的結果。

實驗的時候每個任務都會設置一個兩個對比試驗

• Baseline (FP32): 單精度的激活、權重、梯度得到的模型，運算使用的也是FP32。
• Mixed Precision(MP): FP16用來存儲和數(shù)值運算；權重、激活、梯度都是用的是FP16，其中主備份權重是FP32的。在一些任務中使用了Loss-scaling的技術。運算過程中使用Tensor Cores將累加過程（卷積層、全連接層、矩陣相乘）轉為FP32來計算。

4.1 分類

分類任務上選擇了AlexNet、Vgg-D、GoogLeNet、Inceptionv2、Inceptionv3 和 預激活ResNet50幾個模型進行測試，主要滿足以下條件：

• 使用相同的超參數(shù)比較top-1 準確率。
• 訓練方法采用的是開源庫中默認的方法。
• 數(shù)據(jù)增強方法使用的是最簡單的，并沒有采用開源庫中復雜的方法。主要包括：
• 隨機水平翻轉
• 隨機剪裁crop

可以看到Mixed Precision方法能夠和Baseline方法差不多準確率，有時候甚至會略高于Baseline。

在訓練以上網(wǎng)絡的時候，不需要使用Loss Scale方法，因為這些方法前向傳播和反向傳播的值都在FP16范圍內。

4.2 檢測

這里選擇了Faster RCNN和Multibox-SSD兩種方法在VOC2007數(shù)據(jù)集上進行訓練，兩個方法都是用來VGG-16作為骨干網(wǎng)絡。模型和訓練腳本都來自于開源庫。

可以看到，如果使用Mixed Precision方法，在訓練Multibox SSD的時候可能會由于下溢出導致模型不收斂，但是當使用了Loss Scale技術以后，就可以正常收斂，達到與Baseline相同的結果。

5. 總結

Mixed Precision混合精度是處于一個非常簡單的想法，使用低精度的表示可以節(jié)約顯存、內存同時增加處理器的吞吐量。雖然有以上的種種好處，直接使用半精度會出現(xiàn)一定的問題，比如：下溢出、精度損失等。所以這篇論文核心就是解決使用半精度過程中出現(xiàn)的問題，提出了三個方法達到了非常理想的效果。如果你使用的是有Tensor Core的GPU，那就非常推薦使用混合精度來訓練，只需要安裝NVIDIA提供的Apex庫，然后在你的PyTorch or TensorFlow代碼中加幾行代碼就可以實現(xiàn)。

6. 參考

• https://arxiv.org/pdf/1710.03740v3

• https://blogs.nvidia.com/blog/2019/11/15/whats-the-difference-between-single-double-multi-and-mixed-precision-computing/

• https://zhuanlan.zhihu.com/p/163493798

• https://zhuanlan.zhihu.com/p/79887894