使用Apex进行混合精度训练
使用Apex進行混合精度訓練
轉自：https://fyubang.com/2019/08/26/fp16/
你想獲得雙倍訓練速度的快感嗎？
你想讓你的顯存空間瞬間翻倍嗎？
如果我告訴你只需要三行代碼即可實現(xiàn)，你信不？
在這篇博客里，瓦礫會詳解一下混合精度計算（Mixed Precision），并介紹一款Nvidia開發(fā)的基于PyTorch的混合精度訓練加速神器—Apex，最近Apex更新了API，可以用短短三行代碼就能實現(xiàn)不同程度的混合精度加速，訓練時間直接縮小一半。
話不多說，直接先教你怎么用。
PyTorch實現(xiàn)
from apex import amp model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 這里是“歐一”，不是“零一” with amp.scale_loss(loss, optimizer) as scaled_loss:scaled_loss.backward()對，就是這么簡單，如果你不愿意花時間深入了解，讀到這基本就可以直接使用起來了。
但是如果你希望對FP16和Apex有更深入的了解，或是在使用中遇到了各種不明所以的“Nan”的同學，可以接著讀下去，后面會有一些有趣的理論知識和瓦礫最近一個月使用Apex遇到的各種bug，不過當你深入理解并解決掉這些bug后，你就可以徹底擺脫“慢吞吞”的FP32啦。
理論部分
為了充分理解混合精度的原理，以及API的使用，先補充一點基礎的理論知識。
1. 什么是FP16？
半精度浮點數(shù)是一種計算機使用的二進制浮點數(shù)數(shù)據(jù)類型，使用2字節(jié)（16位）存儲。下圖是FP16和FP32表示的范圍和精度對比。
其中，sign位表示正負，exponent位表示指數(shù)（2n?15+1(n=0)2n?15+1(n=0)），fraction位表示的是分數(shù)（m1024m1024）。其中當指數(shù)為零的時候，下圖加號左邊為0，其他情況為1。下圖是FP16表示范例。
2. 為什么需要FP16？
在使用FP16之前，我想再贅述一下為什么我們使用FP16。
現(xiàn)在模型越來越大，當你使用Bert這一類的預訓練模型時，往往顯存就被模型及模型計算占去大半，當想要使用更大的Batch Size的時候會顯得捉襟見肘。由于FP16的內存占用只有FP32的一半，自然地就可以幫助訓練過程節(jié)省一半的顯存空間。
與普通的空間時間Trade-off的加速方法不同，FP16除了能節(jié)約內存，還能同時節(jié)省模型的訓練時間。在大部分的測試中，基于FP16的加速方法能夠給模型訓練帶來多一倍的加速體驗（爽感類似于兩倍速看肥皂劇）。
硬件的發(fā)展同樣也推動著模型計算的加速，隨著Nvidia張量核心（Tensor Core）的普及，16bit計算也一步步走向成熟，低精度計算也是未來深度學習的一個重要趨勢，再不學習就out啦。
3. FP16帶來的問題：量化誤差
這個部分是整個博客最重要的理論核心。
講了這么多FP16的好處，那么使用FP16的時候有沒有什么問題呢？當然有。FP16帶來的問題主要有兩個：1. 溢出錯誤；2. 舍入誤差。
溢出錯誤（Grad Overflow / Underflow）
由于FP16的動態(tài)范圍（6×10?8～655046×10?8～65504）比FP32的動態(tài)范圍（1.4×10?45～1.7×10381.4×10?45～1.7×1038）要狹窄很多，因此在計算過程中很容易出現(xiàn)上溢出（Overflow，g>65504g>65504）和下溢出（Underflow，g<6×10?8g<6×10?8）的錯誤，溢出之后就會出現(xiàn)“Nan”的問題。
在深度學習中，由于激活函數(shù)的的梯度往往要比權重梯度小，更易出現(xiàn)下溢出的情況。
舍入誤差指的是當梯度過小，小于當前區(qū)間內的最小間隔時，該次梯度更新可能會失敗，用一張圖清晰地表示：
4. 解決問題的辦法：混合精度訓練+動態(tài)損失放大
混合精度訓練的精髓在于“在內存中用FP16做儲存和乘法從而加速計算，用FP32做累加避免舍入誤差”。混合精度訓練的策略有效地緩解了舍入誤差的問題。
即使用了混合精度訓練，還是會存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出（Underflow）。損失放大的思路是：
- 反向傳播前，將損失變化（dLoss）手動增大2k2k倍，因此反向傳播時得到的中間變量（激活函數(shù)梯度）則不會溢出；
- 反向傳播后，將權重梯度縮2k2k倍，恢復正常值。
Apex的新API：Automatic Mixed Precision (AMP)
曾經(jīng)的Apex混合精度訓練的api仍然需要手動half模型已經(jīng)輸入的數(shù)據(jù)，比較麻煩，現(xiàn)在新的api只需要三行代碼即可無痛使用：
from apex import amp model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 這里是“歐一”，不是“零一” with amp.scale_loss(loss, optimizer) as scaled_loss:scaled_loss.backward()opt_level
其中只有一個opt_level需要用戶自行配置：
- O0：純FP32訓練，可以作為accuracy的baseline；
- O1：混合精度訓練（推薦使用），根據(jù)黑白名單自動決定使用FP16（GEMM, 卷積）還是FP32（Softmax）進行計算。
- O2：“幾乎FP16”混合精度訓練，不存在黑白名單，除了Batch norm，幾乎都是用FP16計算。
- O3：純FP16訓練，很不穩(wěn)定，但是可以作為speed的baseline；
動態(tài)損失放大（Dynamic Loss Scaling）
AMP默認使用動態(tài)損失放大，為了充分利用FP16的范圍，緩解舍入誤差，盡量使用最高的放大倍數(shù)（224224），如果產生了上溢出（Overflow），則跳過參數(shù)更新，縮小放大倍數(shù)使其不溢出，在一定步數(shù)后（比如2000步）會再嘗試使用大的scale來充分利用FP16的范圍：
干貨：踩過的那些坑
這一部分是整篇博客最干貨的部分，是瓦礫在最近在apex使用中的踩過的所有的坑，由于apex報錯并不明顯，常常debug得讓人很沮喪，但只要注意到以下的點，95%的情況都可以暢通無阻了：
總結
這篇從理論到實踐地介紹了混合精度計算以及Apex新API（AMP）的使用方法。瓦礫現(xiàn)在在做深度學習模型的時候，幾乎都會第一時間把代碼改成混合精度訓練的了，速度快，精度還不減，確實是調參煉丹必備神器。目前網(wǎng)上還并沒有看到關于AMP以及使用時會遇到的坑的中文博客，所以這一篇也是希望大家在使用的時候可以少花一點時間debug。當然，如果讀者們有發(fā)現(xiàn)新的坑歡迎交流，我會補充在博客中。
Reference
總結
以上是生活随笔為你收集整理的使用Apex进行混合精度训练的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 用防晒霜的防晒功能真的不好吗？
- 下一篇： 一千卡等于多少米