基於Apex的混合精度加速：半精度浮點數FP16

你想獲得雙倍訓練速度的快感嗎？

你想讓你的顯存空間瞬間翻倍嗎？

如果我告訴你只需要三行代碼即可實現，你信不？

在這篇博客里，瓦礫會詳解一下混合精度計算（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的組成

其中，sign位表示正負，exponent位表示指數（  ），fraction位表示的是分數（  ）。其中當指數為零的時候，下圖加號左邊為0，其他情況為1。

fp16的樣例

2. 為什么需要FP16？

在使用FP16之前，我想再贅述一下為什么我們使用FP16。

1. 減少顯存占用 現在模型越來越大，當你使用Bert這一類的預訓練模型時，往往顯存就被模型及模型計算占去大半，當想要使用更大的Batch Size的時候會顯得捉襟見肘。由於FP16的內存占用只有FP32的一半，自然地就可以幫助訓練過程節省一半的顯存空間。
2. 加快訓練和推斷的計算 與普通的空間時間Trade-off的加速方法不同，FP16除了能節約內存，還能同時節省模型的訓練時間。在大部分的測試中，基於FP16的加速方法能夠給模型訓練帶來多一倍的加速體驗（爽感類似於兩倍速看肥皂劇）。
3. 張量核心的普及 硬件的發展同樣也推動着模型計算的加速，隨着Nvidia張量核心（Tensor Core）的普及，16bit計算也一步步走向成熟，低精度計算也是未來深度學習的一個重要趨勢，再不學習就out啦。

3. FP16帶來的問題：量化誤差

這個部分是整個博客最重要的理論核心。 講了這么多FP16的好處，那么使用FP16的時候有沒有什么問題呢？當然有。FP16帶來的問題主要有兩個：1. 溢出錯誤；2. 舍入誤差。

1. 溢出錯誤（Grad Overflow / Underflow） 由於FP16的動態范圍（  ）比FP32的動態范圍（  ）要狹窄很多，因此在計算過程中很容易出現上溢出（Overflow，  ）和下溢出（Underflow，  ）的錯誤，溢出之后就會出現“Nan”的問題。在深度學習中，由於激活函數的的梯度往往要比權重梯度小，更易出現下溢出的情況。

下溢出

2. 舍入誤差（Rounding Error） 舍入誤差指的是當梯度過小，小於當前區間內的最小間隔時，該次梯度更新可能會失敗，用一張圖清晰地表示：

舍入誤差

4. 解決問題的辦法：混合精度訓練+動態損失放大

1. 混合精度訓練（Mixed Precision） 混合精度訓練的精髓在於“在內存中用FP16做儲存和乘法從而加速計算，用FP32做累加避免舍入誤差”。混合精度訓練的策略有效地緩解了舍入誤差的問題。
2. 損失放大（Loss Scaling） 即使用了混合精度訓練，還是會存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出（Underflow）。損失放大的思路是：
1. 反向傳播前，將損失變化（dLoss）手動增大  倍，因此反向傳播時得到的中間變量（激活函數梯度）則不會溢出；
2. 反向傳播后，將權重梯度縮  倍，恢復正常值。

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()

1. opt_level

其中只有一個opt_level需要用戶自行配置：

• O0：純FP32訓練，可以作為accuracy的baseline；
• O1：混合精度訓練（推薦使用），根據黑白名單自動決定使用FP16（GEMM, 卷積）還是FP32（Softmax）進行計算。
• O2：“幾乎FP16”混合精度訓練，不存在黑白名單，除了Batch norm，幾乎都是用FP16計算。
• O3：純FP16訓練，很不穩定，但是可以作為speed的baseline；

2. 動態損失放大（Dynamic Loss Scaling）

AMP默認使用動態損失放大，為了充分利用FP16的范圍，緩解舍入誤差，盡量使用最高的放大倍數（  ），如果產生了上溢出（Overflow），則跳過參數更新，縮小放大倍數使其不溢出，在一定步數后（比如2000步）會再嘗試使用大的scale來充分利用FP16的范圍：

Nvidia關於動態損失放大的解釋

干貨：踩過的那些坑

這一部分是整篇博客最干貨的部分，是瓦礫在最近在apex使用中的踩過的所有的坑，由於apex報錯並不明顯，常常debug得讓人很沮喪，但只要注意到以下的點，95%的情況都可以暢通無阻了：

1. 判斷你的GPU是否支持FP16：支持的有擁有Tensor Core的GPU（2080Ti、Titan、Tesla等），不支持的（Pascal系列）就不建議折騰了。
2. 常數的范圍：為了保證計算不溢出，首先要保證人為設定的常數（包括調用的源碼中的）不溢出，如各種epsilon，INF（根據 
   @喲林小平
    同學的反饋，改成-float('inf')就可以啦）等。
3. Dimension最好是8的倍數：Nvidia官方的文檔的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困擾了我很久。]大誤，參考issue 802 slayton58的說法，注冊好就可以強制加速了。
8. 需要操作模型參數的模塊（類似EMA），要使用AMP封裝后的model。
9. 需要操作梯度的模塊必須在optimizer的step里，不然AMP不能判斷grad是否為Nan。
10. 歡迎補充。。。

總結

這篇從理論到實踐地介紹了混合精度計算以及Apex新API（AMP）的使用方法。瓦礫現在在做深度學習模型的時候，幾乎都會第一時間把代碼改成混合精度訓練的了，速度快，精度還不減，確實是調參煉丹必備神器。目前網上還並沒有看到關於AMP以及使用時會遇到的坑的中文博客，所以這一篇也是希望大家在使用的時候可以少花一點時間debug。當然，如果讀者們有發現新的坑歡迎交流，我會補充在博客中。

 

更優雅的排版請見我的博客：瓦特蘭蒂斯

Reference

1. Intel的低精度表示用於深度學習訓練與推斷
2. Nvidia官方的混合精度訓練文檔
3. Apex官方使用文檔
4. Nvidia-Training Neural Networks with Mixed Precision

編輯於 2020-07-19

原文地址：https://zhuanlan.zhihu.com/p/79887894