混合精度訓練
混合精度訓練是在盡可能減少精度損失的情況下利用半精度浮點數加速訓練。它使用FP16即半精度浮點數存儲權重和梯度。在減少占用內存的同時起到了加速訓練的效果。
IEEE標准中的FP16格式如下:
取值范圍是5.96× 10−8 ~ 65504,而FP32則是1.4×10-45 ~ 3.4×1038。
從FP16的范圍可以看出,用FP16代替原FP32神經網絡計算的最大問題就是精度損失。
float : 1個符號位、8個指數位和23個尾數位
利用fp16 代替 fp32
優點:
1)TensorRT的FP16與FP32相比能有接近一倍的速度提升,前提是GPU支持FP16(如最新的2070,2080,2080ti等)
2)減少顯存。
缺點:
1) 會造成溢出
因此,在日常使用過程中,常使用雙混合精度訓練。如圖:
此過程中的技術:
1) Loss scaling :會存在很多梯度在FP16表達范圍外,我們為了讓其落入半精度范圍內,會給其進行等比放大后縮小。
流程:
對百度&英偉達相關論文的解讀
文中提出了三種避免損失的方法:
1.1. 為每個權重保留一份FP32的副本
在前向和反向時使用FP16,整個過程變成:權重從FP32轉成FP16進行前向計算,得到loss之后,用FP16計算梯度,再轉成FP32更新到FP32的權重上。這里注意得到的loss也是FP32,因為涉及到累加計算(參見下文)。
用FP32保存權重主要是為了避免溢出,FP16無法表示2e-24以下的值,一種是梯度的更新值太小,FP16直接變為了0;二是FP16表示權重的話,和梯度的計算結果也有可能變成0。實驗表明,用FP16保存權重會造成80%的精度損失。
1.2. Loss-scaling
得到FP32的loss后,放大並保存為FP16格式,進行反向傳播,更新時轉為FP32縮放回來。下圖可以看到,很多激活值比較小,無法用FP16表示。因此在前向傳播后對loss進行擴大(固定值或動態值),這樣在反響傳播時所有的值也都擴大了相同的倍數。在更新FP32的權重之前unscale回去。
1.3. 改進算數方法:FP16 * FP16 + FP32。
經過實驗,作者發現將FP16的矩陣相乘后和FP32的矩陣進行加法運算,寫入內存時再轉回FP16可以獲得較好的精度。英偉達V系列GPU卡中的Tensor Core(上圖)也很支持這種操作。因此,在進行大型累加時(batch-norm、softmax),為防止溢出都需要用FP32進行計算,且加法主要被內存帶寬限制,對運算速度不敏感,因此不會降低訓練速度。另外,在進行Point-wise乘法時,受內存帶寬限制。由於算術精度不會影響這些運算的速度,用FP16或者FP32都可以。
2. 實驗結果
從下圖的Accuracy結果可以看到,混合精度基本沒有精度損失:
Loss scale的效果:
3. 如何應用MP
Pytorch可以使用英偉達的開源框架APEX,支持混合進度和分布式訓練:
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
Tensorflow就更簡單了,已經有官方支持,只需要訓練前加一句:
export TF_ENABLE_AUTO_MIXED_PRECISION=1
# 或者
import os
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
參考: