使用Apex進(jìn)行混合精度訓(xùn)練

轉(zhuǎn)自：https://fyubang.com/2019/08/26/fp16/

你想獲得雙倍訓(xùn)練速度的快感嗎？
你想讓你的顯存空間瞬間翻倍嗎？
如果我告訴你只需要三行代碼即可實(shí)現(xiàn)，你信不？

在這篇博客里，瓦礫會(huì)詳解一下混合精度計(jì)算（Mixed Precision），并介紹一款Nvidia開發(fā)的基于PyTorch的混合精度訓(xùn)練加速神器—Apex，最近Apex更新了API，可以用短短三行代碼就能實(shí)現(xiàn)不同程度的混合精度加速，訓(xùn)練時(shí)間直接縮小一半。

話不多說，直接先教你怎么用。

PyTorch實(shí)現(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()

對(duì)，就是這么簡單，如果你不愿意花時(shí)間深入了解，讀到這基本就可以直接使用起來了。

但是如果你希望對(duì)FP16和Apex有更深入的了解，或是在使用中遇到了各種不明所以的“Nan”的同學(xué)，可以接著讀下去，后面會(huì)有一些有趣的理論知識(shí)和瓦礫最近一個(gè)月使用Apex遇到的各種bug，不過當(dāng)你深入理解并解決掉這些bug后，你就可以徹底擺脫“慢吞吞”的FP32啦。

理論部分

為了充分理解混合精度的原理，以及API的使用，先補(bǔ)充一點(diǎn)基礎(chǔ)的理論知識(shí)。

1. 什么是FP16？

半精度浮點(diǎn)數(shù)是一種計(jì)算機(jī)使用的二進(jìn)制浮點(diǎn)數(shù)數(shù)據(jù)類型，使用2字節(jié)（16位）存儲(chǔ)。下圖是FP16和FP32表示的范圍和精度對(duì)比。

其中，sign位表示正負(fù)，exponent位表示指數(shù)（2n?15+1(n=0)2n?15+1(n=0)），fraction位表示的是分?jǐn)?shù)（m1024m1024）。其中當(dāng)指數(shù)為零的時(shí)候，下圖加號(hào)左邊為0，其他情況為1。下圖是FP16表示范例。

2. 為什么需要FP16？

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

減少顯存占用
現(xiàn)在模型越來越大，當(dāng)你使用Bert這一類的預(yù)訓(xùn)練模型時(shí)，往往顯存就被模型及模型計(jì)算占去大半，當(dāng)想要使用更大的Batch Size的時(shí)候會(huì)顯得捉襟見肘。由于FP16的內(nèi)存占用只有FP32的一半，自然地就可以幫助訓(xùn)練過程節(jié)省一半的顯存空間。

加快訓(xùn)練和推斷的計(jì)算
與普通的空間時(shí)間Trade-off的加速方法不同，FP16除了能節(jié)約內(nèi)存，還能同時(shí)節(jié)省模型的訓(xùn)練時(shí)間。在大部分的測試中，基于FP16的加速方法能夠給模型訓(xùn)練帶來多一倍的加速體驗(yàn)（爽感類似于兩倍速看肥皂劇）。

張量核心的普及
硬件的發(fā)展同樣也推動(dòng)著模型計(jì)算的加速，隨著Nvidia張量核心（Tensor Core）的普及，16bit計(jì)算也一步步走向成熟，低精度計(jì)算也是未來深度學(xué)習(xí)的一個(gè)重要趨勢(shì)，再不學(xué)習(xí)就out啦。

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

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

溢出錯(cuò)誤（Grad Overflow / Underflow）
由于FP16的動(dòng)態(tài)范圍（6×10?8～655046×10?8～65504）比FP32的動(dòng)態(tài)范圍（1.4×10?45～1.7×10381.4×10?45～1.7×1038）要狹窄很多，因此在計(jì)算過程中很容易出現(xiàn)上溢出（Overflow，g>65504g>65504）和下溢出（Underflow，g<6×10?8g<6×10?8）的錯(cuò)誤，溢出之后就會(huì)出現(xiàn)“Nan”的問題。

在深度學(xué)習(xí)中，由于激活函數(shù)的的梯度往往要比權(quán)重梯度小，更易出現(xiàn)下溢出的情況。

舍入誤差（Rounding Error）
舍入誤差指的是當(dāng)梯度過小，小于當(dāng)前區(qū)間內(nèi)的最小間隔時(shí)，該次梯度更新可能會(huì)失敗，用一張圖清晰地表示：

4. 解決問題的辦法：混合精度訓(xùn)練+動(dòng)態(tài)損失放大

混合精度訓(xùn)練（Mixed Precision）
混合精度訓(xùn)練的精髓在于“在內(nèi)存中用FP16做儲(chǔ)存和乘法從而加速計(jì)算，用FP32做累加避免舍入誤差”。混合精度訓(xùn)練的策略有效地緩解了舍入誤差的問題。

損失放大（Loss Scaling）
即使用了混合精度訓(xùn)練，還是會(huì)存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出（Underflow）。損失放大的思路是：

• 反向傳播前，將損失變化（dLoss）手動(dòng)增大2k2k倍，因此反向傳播時(shí)得到的中間變量（激活函數(shù)梯度）則不會(huì)溢出；
• 反向傳播后，將權(quán)重梯度縮2k2k倍，恢復(fù)正常值。

Apex的新API：Automatic Mixed Precision (AMP)

曾經(jīng)的Apex混合精度訓(xùn)練的api仍然需要手動(dòng)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

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

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

動(dòng)態(tài)損失放大（Dynamic Loss Scaling）

AMP默認(rèn)使用動(dòng)態(tài)損失放大，為了充分利用FP16的范圍，緩解舍入誤差，盡量使用最高的放大倍數(shù)（224224），如果產(chǎn)生了上溢出（Overflow），則跳過參數(shù)更新，縮小放大倍數(shù)使其不溢出，在一定步數(shù)后（比如2000步）會(huì)再嘗試使用大的scale來充分利用FP16的范圍：

干貨：踩過的那些坑

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

判斷你的GPU是否支持FP16：構(gòu)擁有Tensor Core的GPU（2080Ti、Titan、Tesla等），不支持的（Pascal系列）就不建議折騰了。

常數(shù)的范圍：為了保證計(jì)算不溢出，首先要保證人為設(shè)定的常數(shù)（包括調(diào)用的源碼中的）不溢出，如各種epsilon，INF等。

Dimension最好是8的倍數(shù)：Nvidia官方的文檔的2.2條表示，維度都是8的倍數(shù)的時(shí)候，性能最好。要求維度是8的整數(shù)倍最重要的目的是為了能使用 Tensor Core，Tensor Core 的算力是 CUDA Core 的好多倍。

涉及到sum的操作要小心，很容易溢出，類似Softmax的操作建議用官方API，并定義成layer寫在模型初始化里。

模型書寫要規(guī)范：自定義的Layer寫在模型初始化函數(shù)里，graph計(jì)算寫在forward里。

某些不常用的函數(shù)，在使用前需要注冊(cè)：amp.register_float_function(torch, 'sigmoid')

某些函數(shù)（如einsum）暫不支持FP16加速，建議不要用的太heavy，xlnet的實(shí)現(xiàn)改FP16困擾了我很久。

需要操作模型參數(shù)的模塊（類似EMA），要使用AMP封裝后的model。

需要操作梯度的模塊必須在optimizer的step里，不然AMP不能判斷grad是否為Nan。

歡迎補(bǔ)充。。。

總結(jié)

這篇從理論到實(shí)踐地介紹了混合精度計(jì)算以及Apex新API（AMP）的使用方法。瓦礫現(xiàn)在在做深度學(xué)習(xí)模型的時(shí)候，幾乎都會(huì)第一時(shí)間把代碼改成混合精度訓(xùn)練的了，速度快，精度還不減，確實(shí)是調(diào)參煉丹必備神器。目前網(wǎng)上還并沒有看到關(guān)于AMP以及使用時(shí)會(huì)遇到的坑的中文博客，所以這一篇也是希望大家在使用的時(shí)候可以少花一點(diǎn)時(shí)間debug。當(dāng)然，如果讀者們有發(fā)現(xiàn)新的坑歡迎交流，我會(huì)補(bǔ)充在博客中。

Reference

Intel的低精度表示用于深度學(xué)習(xí)訓(xùn)練與推斷

Nvidia官方的混合精度訓(xùn)練文檔

Apex官方使用文檔

[Nvidia-Training Neural Networks with Mixed Precision](http://on-demand.gputechconf.com/gtc-taiwan/2018/pdf/5-1_Internal Speaker_Michael Carilli_PDF For Sharing.pdf)

總結(jié)

以上是生活随笔為你收集整理的使用Apex进行混合精度训练的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。