1、什么是amp?

amp：Automatic mixed precision，自動混合精度，可以在神經(jīng)網(wǎng)絡推理過程中，針對不同的層，采用不同的數(shù)據(jù)精度進行計算，從而實現(xiàn)節(jié)省顯存和加快速度的目的。

自動混合精度的關鍵詞有兩個：自動、混合精度。

這是由PyTorch 1.6的torch.cuda.amp模塊帶來的：

from torch.cuda import amp

混合精度預示著有不止一種精度的Tensor，那在PyTorch的AMP模塊里是幾種呢？

2種：torch.FloatTensor（浮點型 32位）和torch.HalfTensor（半精度浮點型 16位）；

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

注意

• torch.cuda.amp 的名字意味著這個功能只能在cuda上使用。
• torch默認的tensor精度類型是torch.FloatTensor

2、為什么需要自動混合精度（amp）？

也可以這么問：為什么需要自動混合精度，也就是torch.FloatTensor和torch.HalfTensor的混合，而不全是torch.FloatTensor？或者全是torch.HalfTensor？

原因：

在某些上下文中torch.FloatTensor有優(yōu)勢，在某些上下文中torch.HalfTensor有優(yōu)勢。

torch.HalfTensor

• torch.HalfTensor的優(yōu)勢就是存儲小、計算快、更好的利用CUDA設備的Tensor Core。因此訓練的時候可以減少顯存的占用（可以增加batchsize了），同時訓練速度更快；
• torch.HalfTensor的劣勢就是：數(shù)值范圍?。ǜ菀譕verflow / Underflow）、舍入誤差（Rounding Error，導致一些微小的梯度信息達不到16bit精度的最低分辨率，從而丟失）。

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

• 梯度scale，這正是上一小節(jié)中提到的torch.cuda.amp.GradScaler，通過放大loss的值來防止梯度消失underflow（這只是BP的時候傳遞梯度信息使用，真正更新權(quán)重的時候還是要把放大的梯度再unscale回去）
• 回落到torch.FloatTensor，這就是混合一詞的由來。那怎么知道什么時候用torch.FloatTensor，什么時候用半精度浮點型呢？這是PyTorch框架決定的，AMP上下文中，一些常用的操作中tensor會被自動轉(zhuǎn)化為半精度浮點型的torch.HalfTensor（如：conv1d、conv2d、conv3d、linear、prelu等）

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

答案是 autocast + GradScaler

3.1 autocast

使用torch.cuda.amp模塊中的autocast 類。

from torch.cuda import amp
# 創(chuàng)建model，默認是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 判斷能否使用自動混合精度
enable_amp = True if "cuda" in device.type else False
for input, target in data:
    optimizer.zero_grad()
    # 前向過程(model + loss)開啟 autocast
    with amp.autocast(enabled=enable_amp):
        output = model(input)
        loss = loss_fn(output, target)
    # 反向傳播在autocast上下文之外
    loss.backward()
    optimizer.step()

注意

• 當進入autocast，自動將torch.FloatTensor類型轉(zhuǎn)化為torch.HalfTensor，而不需要手動設置model.half()/input.half，框架會自動做，這也是自動混合精度中“自動”一詞的由來。
• autocast上下文應該只包含網(wǎng)絡的前向過程（包括loss的計算），而不要包含反向傳播。

3.2、GradScaler

這里GradScaler就是第二小節(jié)中提到的梯度scaler模塊，需要在訓練最開始之前使用amp.GradScaler實例化一個GradScaler對象。

from torch.cuda import amp
# 創(chuàng)建model，默認是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 判斷能否使用自動混合精度
enable_amp = True if "cuda" in device.type else False
# 在訓練最開始之前實例化一個GradScaler對象
scaler = amp.GradScaler(enabled=enable_amp)
for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        # 前向過程(model + loss)開啟 autocast
        with amp.autocast(enabled=enable_amp):
            output = model(input)
            loss = loss_fn(output, target)
        # 1、Scales loss.  先將梯度放大 防止梯度消失
        scaler.scale(loss).backward()
        # 2、scaler.step()   再把梯度的值unscale回來.
        # 如果梯度的值不是 infs 或者 NaNs, 那么調(diào)用optimizer.step()來更新權(quán)重,
        # 否則，忽略step調(diào)用，從而保證權(quán)重不更新（不被破壞）
        scaler.step(optimizer)
        # 3、準備著，看是否要增大scaler
        scaler.update()
        # 正常更新權(quán)重
        optimizer.zero_grad()

scaler的大小在每次迭代中動態(tài)的估計，為了盡可能的減少梯度underflow，scaler應該更大；但是如果太大的話，半精度浮點型的tensor又容易overflow（變成inf或者NaN）。

所以動態(tài)估計的原理就是在不出現(xiàn)inf或者NaN梯度值的情況下盡可能的增大scaler的值——在每次scaler.step(optimizer)中，都會檢查是否又inf或NaN的梯度出現(xiàn)：

• 如果出現(xiàn)了inf或者NaN，scaler.step(optimizer)會忽略此次的權(quán)重更新（optimizer.step() )，并且將scaler的大小縮小（乘上backoff_factor）；
• 如果沒有出現(xiàn)inf或者NaN，那么權(quán)重正常更新，并且當連續(xù)多次（growth_interval指定）沒有出現(xiàn)inf或者NaN，則scaler.update()會將scaler的大小增加（乘上growth_factor）。

注意

再強調(diào)一點，amp只能在GPU環(huán)境下使用，因為一來amp是寫在torch.cuda中的函數(shù)，而且amp的中的 amp.GradScaler和amp.autocast函數(shù)構(gòu)造是這樣的：

amp.GradScaler：

    def __init__(self,
                 init_scale=2.**16,
                 growth_factor=2.0,
                 backoff_factor=0.5,
                 growth_interval=2000,
                 enabled=True):
        if enabled and not torch.cuda.is_available():
            warnings.warn("torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.")
            self._enabled = False
        else:
            self._enabled = enabled

amp.autocast：

 def __init__(self, enabled=True):
        if enabled and not torch.cuda.is_available():
            warnings.warn("torch.cuda.amp.autocast only affects CUDA ops, but CUDA is not available.  Disabling.")
            self._enabled = False
        else:
            self._enabled = enabled

4、多GPU訓練

單卡訓練的話上面的代碼已經(jīng)夠了。

要是想多卡跑的話僅僅這樣還不夠，會發(fā)現(xiàn)在forward里面的每個結(jié)果都還是float32的，怎么辦？

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    def forward(self, input_data_c1):
    	with autocast():
    		# code
    	return

只要把model中的forward里面的代碼用autocast代碼塊方式運行就好了。

總結(jié)

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論