【DL】神經網絡混合精度訓練

以前介紹過了神經網絡分佈式訓練，沒怎麼研究混合精度，覺得就是都轉成FP16就行了，最近才發現仍是有些東西的，因此看了下百度和英偉達合做的MIXED PRECISION TRAINING，把細節記下來。

1. 混合精度訓練

混合精度訓練是在儘量減小精度損失的狀況下利用半精度浮點數加速訓練。它使用FP16即半精度浮點數存儲權重和梯度。在減小佔用內存的同時起到了加速訓練的效果。

IEEE標準中的FP16格式以下：

取值範圍是5.96× 10−8 ~ 65504，而FP32則是1.4×10-45 ~ 3.4×1038。

從FP16的範圍能夠看出，用FP16代替原FP32神經網絡計算的最大問題就是精度損失。

文中提出了三種避免損失的方法：

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

放大FP16的梯度，變成FP32的梯度，更新時縮放回來。下圖能夠看到，不少激活值比較小，沒法用FP16表示。所以在前向傳播後對loss進行擴大（固定值或動態值），這樣在反響傳播時全部的值也都擴大了相同的倍數。在更新FP32的權重以前unscale回去。

在訓練期間，一些權重梯度指數很小，在FP16格式下會變成0。爲了解決這個問題，在反向傳播開始時引入一個換算係數，來縮放損失，梯度也經過鏈式法則同時逐漸擴大，並能在FP16中表示出來。

1.3. 改進算數方法：FP16 * FP16 + FP32。

通過實驗，做者發現將FP16的矩陣相乘後和FP32的矩陣進行加法運算，寫入內存時再轉回FP16能夠得到較好的精度。英偉達V系列GPU卡中的Tensor Core（上圖）也很支持這種操做。所以，在進行大型累加時（batch-norm、softmax），爲防止溢出都須要用FP32進行計算，且加法主要被內存帶寬限制，對運算速度不敏感，所以不會下降訓練速度。另外，在進行Point-wise乘法時，用FP16或者FP32均可以，引用原文感覺如下：

Point-wise operations, such as non-linearities and element-wise matrix products, are memory- bandwidth limited. Since arithmetic precision does not impact the speed of these operations, either FP16 or FP32 math can be used.

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'複製代碼

參考：Automatic Mixed Precision for Deep Learning