PyTorch的自動混合精度（AMP）

https://zhuanlan.zhihu.com/p/165152789

 

PyTorch 1.6版本今天發布了，帶來的最大更新就是自動混合精度。release說明的標題是：

1. Stable release of automatic mixed precision (AMP).
2. New Beta features include a TensorPipe backend for RPC, memory profiler,
3. and several improvements to distributed training for both RPC and DDP.

可見自動混合精度正是PyTorch 1.6的最大更新。這就帶來了幾個問題：

1. 什么是自動混合精度訓練？
2. 為什么需要自動混合精度？
3. 如何在PyTorch中使用自動混合精度？

什么是自動混合精度訓練？

我們知道神經網絡框架的計算核心是Tensor，也就是那個從scaler -> array -> matrix -> tensor 維度一路豐富過來的tensor。在PyTorch中，我們可以這樣創建一個Tensor：

>>> import torch >>> gemfield = torch.zeros(70,30) >>> gemfield.type() 'torch.FloatTensor' >>> syszux = torch.Tensor([1,2]) >>> syszux.type() 'torch.FloatTensor'

可以看到默認創建的tensor都是FloatTensor類型。而在PyTorch中，一共有10種類型的tensor：

• torch.FloatTensor (32-bit floating point)
• torch.DoubleTensor (64-bit floating point)
• torch.HalfTensor (16-bit floating point 1)
• torch.BFloat16Tensor (16-bit floating point 2)
• torch.ByteTensor (8-bit integer (unsigned))
• torch.CharTensor (8-bit integer (signed))
• torch.ShortTensor (16-bit integer (signed))
• torch.IntTensor (32-bit integer (signed))
• torch.LongTensor (64-bit integer (signed))
• torch.BoolTensor (Boolean)

由此可見，默認的Tensor是32-bit floating point，這就是32位浮點型精度的Tensor。

自動混合精度的關鍵詞有兩個：自動、混合精度。這是由PyTorch 1.6的torch.cuda.amp模塊帶來的：

from torch.cuda.amp import autocast as autocast

混合精度預示着有不止一種精度的Tensor，那在PyTorch的AMP模塊里是幾種呢？2種：torch.FloatTensor和torch.HalfTensor；

自動預示着Tensor的dtype類型會自動變化，也就是框架按需自動調整tensor的dtype（其實不是完全自動，有些地方還是需要手工干預）；

torch.cuda.amp 的名字意味着這個功能只能在cuda上使用，事實上，這個功能正是NVIDIA的開發人員貢獻到PyTorch項目中的。而只有支持Tensor core的CUDA硬件才能享受到AMP的好處（比如2080ti顯卡）。Tensor Core是一種矩陣乘累加的計算單元，每個Tensor Core每個時鍾執行64個浮點混合精度操作（FP16矩陣相乘和FP32累加），英偉達宣稱使用Tensor Core進行矩陣運算可以輕易的提速，同時降低一半的顯存訪問和存儲。

因此，在PyTorch中，當我們提到自動混合精度訓練，我們說的就是在NVIDIA的支持Tensor core的CUDA設備上使用torch.cuda.amp.autocast （以及torch.cuda.amp.GradScaler）來進行訓練。咦？為什么還要有torch.cuda.amp.GradScaler？

為什么需要自動混合精度？

這個問題其實暗含着這樣的意思：為什么需要自動混合精度，也就是torch.FloatTensor和torch.HalfTensor的混合，而不全是torch.FloatTensor？或者全是torch.HalfTensor？

如果非要以這種方式問，那么答案只能是，在某些上下文中torch.FloatTensor有優勢，在某些上下文中torch.HalfTensor有優勢唄。答案進一步可以轉化為，相比於之前的默認的torch.FloatTensor，torch.HalfTensor有時具有優勢，有時劣勢不可忽視。

torch.HalfTensor的優勢就是存儲小、計算快、更好的利用CUDA設備的Tensor Core。因此訓練的時候可以減少顯存的占用（可以增加batchsize了），同時訓練速度更快；

torch.HalfTensor的劣勢就是：數值范圍小（更容易Overflow / Underflow）、舍入誤差（Rounding Error，導致一些微小的梯度信息達不到16bit精度的最低分辨率，從而丟失）。

可見，當有優勢的時候就用torch.HalfTensor，而為了消除torch.HalfTensor的劣勢，我們帶來了兩種解決方案：

1，梯度scale，這正是上一小節中提到的torch.cuda.amp.GradScaler，通過放大loss的值來防止梯度的underflow（這只是BP的時候傳遞梯度信息使用，真正更新權重的時候還是要把放大的梯度再unscale回去）；

2，回落到torch.FloatTensor，這就是混合一詞的由來。那怎么知道什么時候用torch.FloatTensor，什么時候用半精度浮點型呢？這是PyTorch框架決定的，在PyTorch 1.6的AMP上下文中，如下操作中tensor會被自動轉化為半精度浮點型的torch.HalfTensor：

1. __matmul__
2. addbmm
3. addmm
4. addmv
5. addr
6. baddbmm
7. bmm
8. chain_matmul
9. conv1d
10. conv2d
11. conv3d
12. conv_transpose1d
13. conv_transpose2d
14. conv_transpose3d
15. linear
16. matmul
17. mm
18. mv
19. prelu

如何在PyTorch中使用自動混合精度？

答案就是autocast + GradScaler。

1，autocast

正如前文所說，需要使用torch.cuda.amp模塊中的autocast 類。使用也是非常簡單的：

from torch.cuda.amp import autocast as autocast

# 創建model，默認是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # 前向過程(model + loss)開啟 autocast
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # 反向傳播在autocast上下文之外
    loss.backward()
    optimizer.step()

可以使用autocast的context managers語義（如上所示），也可以使用decorators語義。 當進入autocast的上下文后，上面列出來的那些CUDA ops 會把tensor的dtype轉換為半精度浮點型，從而在不損失訓練精度的情況下加快運算。剛進入autocast的上下文時，tensor可以是任何類型，你不要在model或者input上手工調用.half() ，框架會自動做，這也是自動混合精度中“自動”一詞的由來。

另外一點就是，autocast上下文應該只包含網絡的前向過程（包括loss的計算），而不要包含反向傳播，因為BP的op會使用和前向op相同的類型。

還有的時候呀，你的代碼在autocast上下文中會報如下的錯誤：

Traceback (most recent call last):
......
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
......
RuntimeError: expected scalar type float but found c10::Half

對於RuntimeError: expected scalar type float but found c10::Half，這估計是個bug。你可以在tensor上手工調用.float()來讓type匹配。

2，GradScaler

但是別忘了前面提到的梯度scaler模塊呀，需要在訓練最開始之前實例化一個GradScaler對象。因此PyTorch中經典的AMP使用方式如下：

from torch.cuda.amp import autocast as autocast

# 創建model，默認是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在訓練最開始之前實例化一個GradScaler對象
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # 前向過程(model + loss)開啟 autocast
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss. 為了梯度放大.
        scaler.scale(loss).backward()

        # scaler.step() 首先把梯度的值unscale回來.
        # 如果梯度的值不是 infs 或者 NaNs, 那么調用optimizer.step()來更新權重,
        # 否則，忽略step調用，從而保證權重不更新（不被破壞）
        scaler.step(optimizer)

        # 准備着，看是否要增大scaler
        scaler.update()

scaler的大小在每次迭代中動態的估計，為了盡可能的減少梯度underflow，scaler應該更大；但是如果太大的話，半精度浮點型的tensor又容易overflow（變成inf或者NaN）。所以動態估計的原理就是在不出現inf或者NaN梯度值的情況下盡可能的增大scaler的值——在每次scaler.step(optimizer)中，都會檢查是否又inf或NaN的梯度出現：

1，如果出現了inf或者NaN，scaler.step(optimizer)會忽略此次的權重更新（optimizer.step() )，並且將scaler的大小縮小（乘上backoff_factor）；

2，如果沒有出現inf或者NaN，那么權重正常更新，並且當連續多次（growth_interval指定）沒有出現inf或者NaN，則scaler.update()會將scaler的大小增加（乘上growth_factor）。

最后

你也可以使用我們提供的PyTorch項目規范來簡化開發：

https://github.com/deepVAC/deepvac/​github.com

繼承自DeepvacTrain類，在deepvac_config中設置config.amp = True即可。

在Gemfield的一個Conv2d和全連接層占主導的網絡中，當開啟AMP后，訓練時顯存的占用從11GB下降到了8GB。而速度......