混合精度訓練-Pytorch

1、需求解讀

  作爲一名算法工程師，我們經常會遇到訓練網絡的事情，當前訓練網絡的整個過程基本上都是在N卡上面執行的，當我們的數據集比較大時，訓練網絡會耗費大量的時間。由於我們需要使用反向傳播來更新具有細微變化的權重，因而我們在訓練網絡的過程中通常會選用FP32類型的數據和權重。說了這麼多，那麼混合精度到底是什麼呢，有什麼用呢？
  簡而言之，所謂的混合精度訓練，即當你使用N卡訓練你的網絡時，混合精度會在內存中用FP16做儲存和乘法從而加速計算，用FP32做累加避免舍入誤差。它的優勢就是可以使你的訓練時間減少一半左右。它的缺陷是隻能在支持FP16操作的一些特定類型的顯卡上面使用，而且會存在溢出誤差和舍入誤差。

2、F16和FP32的區別與聯繫

聯繫：

• FP32和FP16都是用來表示某一個數值；
• FP32和FP16都是由符號位、指數和尾數一起組成；

區別：

• FP32由1位符號位、8位指數和23位尾數組成，FP16由1位符號位、5位指數和10位尾數組成；
• FP32能夠表示的範圍是$1.4 \times 10^{-45}<x<3.4 \times 10^{38}$，FP16能夠表示的範圍是$5.96 \times 10^{-8}<x<65504$。即FP16最大能夠表示的數字是65503；
• FP32能夠更加準確的表示某一個數字；

3、F16優點簡介

優點1-FP16計算速度更快、更加節約內存

  上圖展示了FP16和FP32在內存消耗上面的不同之處。通過觀察上圖我們可以得出：

• 計算同樣的操作，FP16可以獲得8倍的加速；
• FP16能夠獲得2倍左右的內存扇出；
• FP16能夠節省1/2的內存資源；

優點2-FP16可以使用上特定顯卡中專門爲加速所設計的Tensor Core

  上圖展示了執行卷積的過程（乘操作和加操作）。使用FP16執行成操作，然後使用FP16或者FP32執行乘操作。與使用FP32計算相比，在Volta V100（該架構中存在Tensor Core，支持FP16操作）上面可以獲得8倍的性能提速，最終達到125TFlops的扇出。

4、F16缺點簡介

缺點1-FP16會帶來梯度溢出錯誤
  Grad Overflow / Underflow，即梯度溢出。由於FP16的動態範圍是$5.96 \times 10^{-8}<x<65504$，比FP32的動態範圍小了很多，因而在計算的過程中很容易出現上溢出（超出能夠表示的最大數值）和下溢出（超出能夠表示的最小數值）問題，溢出之後就會出現NAN的問題。在深度學習中，由於激活函數的的梯度往往要比權重梯度小，更易出現下溢出的情況，具體的細節如下圖所示。

缺點2-FP16會帶來舍入誤差
  舍入誤差，即當梯度過小，小於當前區間內的最小間隔時，該次梯度更新可能會失敗，具體的細節如下圖所示，由於更新的梯度值超出了FP16能夠表示的最小值的範圍，因此該數值將會被捨棄，這個權重將不進行更新。

解決方案：

• 使用混合精度訓練。所謂的混合精度訓練，即在內存中用FP16做儲存和乘法從而加速計算，用FP32做累加避免舍入誤差，這樣可以很好的解決舍入誤差的問題。
• 損失放大。有些情況下，即使使用了混合精度訓練的方法，由於激活梯度的值太小，會造成下溢出，從而導致模型無法收斂的問題。所謂的損失放大，即反向傳播前，將損失變化（dLoss）手動增大$2^{k}$倍，因此反向傳播時得到的中間變量（激活函數梯度）則不會溢出；反向傳播後，將權重梯度縮$2^{k}$倍，恢復正常值。

5、混合精度訓練代碼實戰

5.1 代碼實現

使用FP32訓練代碼如下所示：

# coding=utf-8
import torch

N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device=“cuda”)
y = torch.randn(N, D_out, device=“cuda”)
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for t in range(500):
	y_pred = model(x)
	loss = torch.nn.functional.mse_loss(y_pred, y)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()

使用FP16訓練代碼如下所示，僅僅需要在原始的Pytorch代碼中增加3行代碼，你就可以體驗到極致的性能加速啦。

# coding=utf-8
import torch

N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device=“cuda”)
y = torch.randn(N, D_out, device=“cuda”)
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
model, optimizer = amp.initialize(model, optimizer, opt_level=“O1”)
for t in range(500):
	y_pred = model(x)
	loss = torch.nn.functional.mse_loss(y_pred, y)
	optimizer.zero_grad()
	with amp.scale_loss(loss, optimizer) as scaled_loss:
		scaled_loss.backward()
	optimizer.step()

5.2 代碼解析

1、model, optimizer = amp.initialize(model, optimizer, opt_level=“O1”)

  這行代碼的主要作用是對模型和優化器執行初始化操作，方便後續的混合精度訓練。其中opt_level表示優化的等級，當前支持4個等級的優化，具體的情況如下圖所示：

• 當opt_level='00’時，表示的是當前執行FP32訓練，即正常的訓練，當前優化等級執行的具體操作是cast_model_type=torch.float32、patch_torch_function= False、keep_batchnorm_fp32=None、master_weight=False、loss_scale=1.0。
• 當opt_level='01’時，表示的是當前使用部分FP16混合訓練，當前優化等級執行的具體操作是cast_model_type=None、patch_torch_function=True、keep_batch norm_fp32=None、master_weight=None、loss_scale=“dynamic”。
  
• 當opt_level='02’時，表示的是除了BN層的權重外，其他層的權重都使用FP16執行訓練，當前優化等級執行的具體操作是cast_model_type=torch.float16、patch _torch_function=False、keep_batchnorm_fp32=True、master_weight =True 、loss_scale=“dynamic”。
• 當opt_level='03’時，表示的是默認所有的層都使用FP16執行計算，當keep_batch norm_fp32=True，則會使用cudnn執行BN層的計算，該優化等級能夠獲得最快的速度，但是精度可能會有一些較大的損失。當前優化等級執行的具體操作是cast_ model_type=torch.float16、patch _torch_function=False、keep_batchnorm _fp32=False、master_weight =False、loss_scale=1.0

注意事項：
1、cast_model_type表示的是模型的輸入類型。當前支持的類型包括torch.float32和torch.float16；
2、patch _torch_function表示的是根據不同函數的輸入數據要求獲得一個最優的輸入類型。GEMM和Convolution等運算可以使用FP16快速的獲得最終的結果，由於softma x/exponentiation/pow等運算需要較高的精度，所以選擇使用FP32來計算。當前支持的類型包括False和True。
3、keep_batchnorm _fp32表示的是是否需要對網絡中BN層執行特殊處理。由於網絡中的BN層會影響數據的分佈情況，從而進一步影響網絡的訓練過程，因此需要認真的去處理這個類型的層。當該層使用FP32時，網絡的訓練過程會更加穩定。當前支持的類型包括False和True。

4、master_weight 表示的是網絡在訓練過程中部分參數使用FP32來表示，部分參數使用FP16來表示。上圖中藍色的框表示FP32類型，綠色的框表示FP16類型，FP32在轉化爲FP16的過程中會進行備份（master_0）,optimizer都是使用FP32來表示，而model部分中部分參數是FP16類型，部分參數是FP32類型，梯度更新的過程通常是在master上面執行。
5、loss_scale表示是是否需要執行損失放大操作。1.0表示不需要執行損失放大操作，dynamic表示需要執行損失放大操作。

NVIDIA官方給出的使用規則如下所示：

• 首先，建議將opt_level設置爲00。即使用FP32訓練模型，從而建立起一個準確的Baseline;
• 然後，嘗試着將opt_level設置爲01。即嘗試着使用混合精度訓練方法；
• 接着，如果你對訓練的速度有着較高的要求，建議將opt_level設置爲02或者03；

2、 with amp.scale_loss(loss, optimizer) as scaled_loss:
      scaled_loss.backward()
  這行代碼的主要作用是在反向傳播前進行梯度放大來進行更新，在反向傳播後進行梯度縮放，返回原來的值，但是可以很好的解決由於梯度值太小模型無法更新的問題。具體的細節如下圖所示：

  上圖展示了FP16在計算的過程中由於梯度值太小，超出了FP16能表示的下限值，因而無法進行權重更新，導致網絡不收斂。

  上圖展示了使用損失方法（Scaled Loss）的方法來很好的解決這個問題，即在反向傳播之前，給這些比較少的數值乘上$2^{k}$，即將其擴大$2^{k}$倍，將其調整到FP16能夠支持的一個合理的範圍內，那麼FP16就可以對這個比較小的梯度增量值執行更新，這樣就可以很好的解決這個問題。

  上圖展示了對反向傳播之後的結果之後後處理的過程，由於我們爲了解決方向傳播之前梯度數值太小而將它擴大$2^{k}$倍，那麼這樣計算之後就相當於我們認爲的將梯度值增加了$2^{k}$倍，爲了獲得準確的權重值，我們需要在反向傳播之後除以$2^{k}$，整個過程在optimizer.step()執行之前。

6、F16訓練效果展示

  上圖展示了使用混合訓練在多個經典模型上面的加速效果。在BERT模型上，使用混合精度訓練可以獲得4倍的提速。換句話說，我們原本需要4天才能訓練好的模型，現在1天就可以訓練出來，而且能夠達到幾乎相同的精度級別。這在很多情況下還是挺有用的，這個方法在減少模型訓練時間的同時可以節省更多的電費，除此之外，可以節約算法工程師們的時間，從而提高他們的工作效率。

  上圖展示了使用混合精度訓練的模型的精度。通過觀察我們可以得出以下的結論：混合精度訓練在提升訓練速度的同時可以達到和FP32訓練同樣的精度。

7、個人總結

  通過仔細理解上面的內容，你應該會對混合精度訓練有了一個全新的認識。所謂的混合精度訓練，即當你使用N卡訓練你的網絡時，混合精度會在內存中用FP16做儲存和乘法從而加速計算，用FP32做累加避免舍入誤差。它的優勢就是可以使你的訓練時間減少一半左右。它的缺陷是隻能在支持FP16操作的一些特定類型的顯卡上面使用，而且會存在溢出誤差和舍入誤差。總而言之，混合精度訓練可以在保證精度的同時極大的提升你的訓練速度，如果你習慣使用pytorch來訓練網絡，那你就可以獲得極致的訓練速度啦。當前混合精度訓練仍然存在着一些限制條件，首先，你的硬件設備需要支持FP16計算；其次，你的硬件設備需要具有Tensor_Core單元（這僅僅存在於一些新架構的N卡上面）；接着，當前的僅有少量的深度學習框架支持混合精度訓練（Pytorch）;最後，混合精度不僅僅可以用在網絡訓練的過程中，同樣也可以將它應用在網絡推理過程中執行加速。隨着越來越多的硬件設備支持FP16計算之後，混合精度訓練和推理應該會成爲一個首選，我相信越來越多的訓練和推理框架都會在短期內逐漸支持混合精度訓練。

參考資料

[1] 參考博客
[2] NVIDIA參考資料
[3] GTC_2019
[4] apex

注意事項

[1] 該博客是本人原創博客，如果您對該博客感興趣，想要轉載該博客，請與我聯繫（qq郵箱：[email protected]）,我會在第一時間回覆大家，謝謝大家的關注.
[2] 由於個人能力有限，該博客可能存在很多的問題，希望大家能夠提出改進意見。
[3] 如果您在閱讀本博客時遇到不理解的地方，希望您可以聯繫我，我會及時的回覆您，和您交流想法和意見，謝謝。
[4] 本人業餘時間承接各種本科畢設設計和各種小項目，包括圖像處理（數據挖掘、機器學習、深度學習等）、matlab仿真、python算法及仿真等，有需要的請加QQ：1575262785詳聊，備註“項目”！！！