作者丨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的混合精度加速的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。