PyTorch必备神器 | 唯快不破：基于Apex的混合精度加速
作者丨Nicolas
單位丨追一科技AI Lab研究員
研究方向丨信息抽取、機器閱讀理解
你想獲得雙倍訓練速度的快感嗎？?
你想讓你的顯卡內存瞬間翻倍嗎？?
如果告訴你只需要三行代碼即可實現，你信不？?
在這篇文章里，筆者會詳解一下混合精度計算（Mixed Precision），并介紹一款 NVIDIA 開發的基于 PyTorch 的混合精度訓練加速神器——Apex，最近 Apex 更新了 API，可以用短短三行代碼就能實現不同程度的混合精度加速，訓練時間直接縮小一半。?
話不多說，直接先教你怎么用。
PyTorch實現
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？
半精度浮點數是一種計算機使用的二進制浮點數數據類型，使用?2?字節（16?位）存儲。
▲?FP16和FP32表示的范圍和精度對比
?
其中，?sign?位表示正負，?exponent?位表示指數，?fraction?位表示的是分數。其中當指數為零的時候，下圖加號左邊為 0，其他情況為?1。
▲?FP16的表示范例
?
2.?為什么需要FP16？
在使用?FP16?之前，我想再贅述一下為什么我們使用?FP16。
減少顯存占用?現在模型越來越大，當你使用?Bert?這一類的預訓練模型時，往往模型及模型計算就占去顯存的大半，當想要使用更大的?Batch Size?的時候會顯得捉襟見肘。由于 FP16?的內存占用只有?FP32?的一半，自然地就可以幫助訓練過程節省一半的顯存空間。
加快訓練和推斷的計算?與普通的空間時間?Trade-off?的加速方法不同，FP16?除了能節約內存，還能同時節省模型的訓練時間。在大部分的測試中，基于 FP16?的加速方法能夠給模型訓練帶來多一倍的加速體驗（爽感類似于兩倍速看肥皂劇）。
張量核心的普及?硬件的發展同樣也推動著模型計算的加速，隨著?NVIDIA?張量核心（Tensor Core）的普及，16bit?計算也一步步走向成熟，低精度計算也是未來深度學習的一個重要趨勢，再不學習就?out?啦。
?
3. FP16帶來的問題：量化誤差
這個部分是整篇文章最重要的理論核心。
?
講了這么多?FP16?的好處，那么使用?FP16?的時候有沒有什么問題呢？當然有。FP16 帶來的問題主要有兩個：1.?溢出錯誤；2.?舍入誤差。
?
溢出錯誤（Grad Overflow / Underflow）由于?FP16?的動態范圍比?FP32?的動態范圍要狹窄很多，因此在計算過程中很容易出現上溢出（Overflow，g>65504）和下溢出（Underflow，）的錯誤，溢出之后就會出現“Nan”的問題。
在深度學習中，由于激活函數的的梯度往往要比權重梯度小，更易出現下溢出的情況。
▲?下溢出問題
舍入誤差（Rounding Error）舍入誤差指的是當梯度過小，小于當前區間內的最小間隔時，該次梯度更新可能會失敗，用一張圖清晰地表示：
▲?舍入誤差
?
4.?解決問題的辦法：混合精度訓練+動態損失放大
混合精度訓練（Mixed Precision）混合精度訓練的精髓在于“在內存中用?FP16?做儲存和乘法從而加速計算，用?FP32?做累加避免舍入誤差”。混合精度訓練的策略有效地緩解了舍入誤差的問題。
損失放大（Loss Scaling）即使用了混合精度訓練，還是會存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出（Underflow）。損失放大的思路是：
反向傳播前，將損失變化（dLoss）手動增大倍，因此反向傳播時得到的中間變量（激活函數梯度）則不會溢出；
反向傳播后，將權重梯度縮小倍，恢復正常值。
Apex的新API：Automatic Mixed Precision (AMP)
曾經的 Apex 混合精度訓練的 API 仍然需要手動 half 模型以及輸入的數據，比較麻煩，現在新的 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?：混合精度訓練（推薦使用），根據黑白名單自動決定使用 FP16（GEMM, 卷積）還是 FP32（Softmax）進行計算；?
?O2 ：“幾乎 FP16”混合精度訓練，不存在黑白名單，除了 Batch Norm，幾乎都是用 FP16 計算；
?O3 ：純 FP16 訓練，很不穩定，但是可以作為 speed 的 baseline。
動態損失放大（Dynamic Loss Scaling）?
AMP 默認使用動態損失放大，為了充分利用 FP16 的范圍，緩解舍入誤差，盡量使用最高的放大倍數（），如果產生了上溢出（Overflow），則跳過參數更新，縮小放大倍數使其不溢出，在一定步數后（比如 2000 步）會再嘗試使用大的 scale 來充分利用 FP16 的范圍：
▲?AMP中動態損失放大的策略
干貨：踩過的那些坑
這一部分是整篇文章最干貨的部分，是筆者在最近在 apex 使用中的踩過的所有的坑，由于 apex 報錯并不明顯，常常 debug 得讓人很沮喪，但只要注意到以下的點，95% 的情況都可以暢通無阻了：?
1. 判斷你的 GPU 是否支持 FP16：支持的有擁有 Tensor Core 的 GPU（2080Ti、Titan、Tesla 等），不支持的（Pascal 系列）就不建議折騰了；
2. 常數的范圍：為了保證計算不溢出，首先要保證人為設定的常數（包括調用的源碼中的）不溢出，如各種 epsilon，INF 等；
3. Dimension 最好是 8 的倍數：NVIDIA 官方的文檔的 2.2 條 [2] 表示，維度都是 8 的倍數的時候，性能最好；
4. 涉及到 sum 的操作要小心，很容易溢出，類似 Softmax 的操作建議用官方 API，并定義成 layer 寫在模型初始化里；
5. 模型書寫要規范：自定義的 Layer 寫在模型初始化函數里，graph 計算寫在 forward 里；
6. 某些不常用的函數，在使用前需要注冊：?amp.register_float_function(torch, 'sigmoid') ；
7. 某些函數（如 einsum）暫不支持 FP16 加速，建議不要用的太 heavy，XLNet 的實現改 FP16 [4] 困擾了我很久；
8. 需要操作模型參數的模塊（類似 EMA），要使用 AMP 封裝后的 model；
9. 需要操作梯度的模塊必須在 optimizer 的 step 里，不然 AMP 不能判斷 grad 是否為 Nan；
10. 歡迎補充。
總結
這篇從理論到實踐地介紹了混合精度計算以及 Apex 新 API（AMP）的使用方法。筆者現在在做深度學習模型的時候，幾乎都會第一時間把代碼改成混合精度訓練的了，速度快，精度還不減，確實是調參煉丹必備神器。目前網上還并沒有看到關于 AMP 以及使用時會遇到的坑的中文博客，所以這一篇也是希望大家在使用的時候可以少花一點時間 debug。當然，如果讀者們有發現新的坑歡迎交流，我會補充在專欄的博客 [5] 中。
Reference
[1] Intel的低精度表示用于深度學習訓練與推斷?
http://market.itcgb.com/Contents/Intel/OR_AI_BJ/images/Brian_DeepLearning_LowNumericalPrecision.pdf
[2] NVIDIA官方混合精度訓練文檔
https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html?
[3] Apex官方使用文檔
https://nvidia.github.io/apex/amp.html
[4]?XLNet的實現改FP16
https://github.com/NVIDIA/apex/issues/394?
[5] 專欄博客
https://zhuanlan.zhihu.com/p/79887894
點擊以下標題查看更多往期內容：?
ICCV 2019 | 基于持續學習的條件圖像生成模型
@即將開學的你，請收好這份必讀論文清單
小米AutoML團隊發布可伸縮超網SCARLET
Github大熱論文 | 基于GAN的新型無監督圖像轉換
后BERT時代的那些NLP預訓練模型
KDD Cup 2019 AutoML Track冠軍團隊技術分享
#投 稿 通 道#
?讓你的論文被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。
??來稿標準：
? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?
? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?
? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志
? 投稿郵箱：
? 投稿郵箱：hr@paperweekly.site?
? 所有文章配圖，請單獨在附件中發送?
? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通
?
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
關于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
▽ 點擊 |?閱讀原文?| 獲取最新論文推薦
總結
以上是生活随笔為你收集整理的PyTorch必备神器 | 唯快不破：基于Apex的混合精度加速的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 北京 | 免费高效训练及OpenVINO
- 下一篇： “让Keras更酷一些！”：层中层与ma