天天看點
傳回

Pytorch自動混合精度(AMP)介紹與使用Pytorch自動混合精度(AMP)介紹與使用

Pytorch自動混合精度(AMP)介紹與使用

背景：

pytorch從1.6版本開始，已經内置了torch.cuda.amp，采用自動混合精度訓練就不需要加載第三方NVIDIA的apex庫了。本文主要從三個方面來介紹AMP：

一．什麼是AMP?

二．為什麼要使用AMP？

三．如何使用AMP?

四. 注意事項

正文：

一．什麼是AMP?

預設情況下，大多數深度學習架構都采用32位浮點算法進行訓練。２０１７年，NVIDIA研究了一種用于混合精度訓練的方法，該方法在訓練網絡時将單精度（FP32）與半精度(FP16)結合在一起，并使用相同的超參數實作了與FP32幾乎相同的精度。

在介紹AMP之前，先來了解下FP16與FP32，FP16也即半精度是一種計算機使用的二進制浮點資料類型，使用2位元組存儲。而FLOAT就是FP32。

Pytorch自動混合精度(AMP)介紹與使用Pytorch自動混合精度(AMP)介紹與使用

其中，sign位表示正負，exponent位表示指數2^(n-15+1(n=0))，fraction位表示分數(m/1024)。

Pytorch自動混合精度(AMP)介紹與使用Pytorch自動混合精度(AMP)介紹與使用

一般情況下，我們在pytorch中建立一個Tensor:

>>import torch
>>tensor1=torch.zeros(30,20)
>>tensor1.type()
'torch.FloatTensor'

>>tensor2=torch.Tensor([1,2])
>>tensor2.type()

'torch.FlatTensor'

可以看到，預設建立的tensor都是FloatTensor類型。而在Pytorch中，一共有10種類型的tensor:

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

預設Tensor是32bit floating point，這就是32位浮點型精度的tensor。

AMP(自動混合精度）的關鍵詞有兩個：自動，混合精度。

自動：Tensor的dtype類型會自動變化，架構按需自動調整tensor的dtype,當然有些地方還需手動幹預。

混合精度：采用不止一種精度的Tensor，torch.FloatTensor和torch.HalfTensor

pytorch1.6的新包：torch.cuda.amp，是ＮVIDIA開發人員貢獻到pytorch裡的。隻有支援tensor core的CUDA硬體才能享受到AMP帶來的優勢。Tensor core是一種矩陣乘累加的計算單元，每個tensor core時針執行64個浮點混合精度操作（FP16矩陣相乘和FP32累加）。

二、為什麼要使用AMP?

​ 前面已介紹，AMP其實就是Float32與Float16的混合，那為什麼不單獨使用Float32或Float16，而是兩種類型混合呢？原因是：在某些情況下Float32有優勢，而在另外一些情況下Float16有優勢。這裡先介紹下FP16：

優勢有三個：

１．減少顯存占用；

２．加快訓練和推斷的計算，能帶來多一倍速的體驗；

３．張量核心的普及（NVIDIA　Tensor Core）,低精度計算是未來深度學習的一個重要趨勢。

但凡事都有兩面性，FP16也帶來了些問題：１．溢出錯誤；２．舍入誤差；

１．溢出錯誤：由于FP16的動态範圍比FP32位的狹窄很多，是以，在計算過程中很容易出現上溢出和下溢出，溢出之後就會出現"NaN"的問題。在深度學習中，由于激活函數的梯度往往要比權重梯度小，更易出現下溢出的情況

Pytorch自動混合精度(AMP)介紹與使用Pytorch自動混合精度(AMP)介紹與使用

2.舍入誤差

舍入誤差指的是當梯度過小時，小于目前區間内的最小間隔時，該次梯度更新可能會失敗：

Pytorch自動混合精度(AMP)介紹與使用Pytorch自動混合精度(AMP)介紹與使用

為了消除torch.HalfTensor也就是FP16的問題，需要使用以下兩種方法：

１）混合精度訓練

在記憶體中用FP16做儲存和乘法進而加速計算，而用FP32做累加避免舍入誤差。混合精度訓練的政策有效地緩解了舍入誤差的問題。

什麼時候用torch.FloatTensor,什麼時候用torch.HalfTensor呢？這是由pytorch架構決定的，在pytorch1.6的AMP上下文中，以下操作中Tensor會被自動轉化為半精度浮點型torch.HalfTensor：

__matmul__
addbmm
addmm
addmv
addr
baddbmm
bmm
chain_matmul
conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear
matmul
mm
mv
prelu

2)損失放大（Loss scaling)

​ 即使了混合精度訓練，還是存在無法收斂的情況，原因是激活梯度的值太小，造成了溢出。可以通過使用torch.cuda.amp.GradScaler，通過放大loss的值來防止梯度的underflow（隻在BP時傳遞梯度資訊使用，真正更新權重時還是要把放大的梯度再unscale回去）；

反向傳播前，将損失變化手動增大2^k倍，是以反向傳播時得到的中間變量（激活函數梯度）則不會溢出；

反向傳播後，将權重梯度縮小2^k倍，恢複正常值。

三．如何使用AMP?

目前有兩種版本：pytorch1.5之前使用的NVIDIA的三方包apex.amp和pytorch1.6自帶的torch.cuda.amp

１．pytorch1.5之前的版本（包括１．５）

使用方法如下：

from apex import amp
model,optimizer = amp.initial(model,optimizer,opt_level="O1")   #注意是O,不是０
with amp.scale_loss(loss,optimizer) as scaled_loss:
    scaled_loss.backward()取代loss.backward()

其中,opt_level配置如下：

O0:純FP32訓練，可作為accuracy的baseline；

O1:混合精度訓練（推薦使用），根據黑白名單自動決定使用FP16(GEMM,卷積）還是FP32（softmax)進行計算。

O2:幾乎FP16，混合精度訓練，不存在黑白名單　，除了bacthnorm，幾乎都是用FP16計算；

O3:純FP16訓練，很不穩定，但是可以作為speed的baseline；

動态損失放大（dynamic loss scaling)部分，為了充分利用FP16的範圍，緩解舍入誤差，盡量使用最高的放大倍數2^24,如果産生上溢出，則跳出參數更新，縮小放大倍數使其不溢出。在一定步數後再嘗試使用大的scale來充分利用FP16的範圍。

分布式訓練：

import argparse
import apex import amp
import apex.parallel import convert_syncbn_model
import apex.parallel import DistributedDataParallel as DDP

定義超參數：
def parse():
    parser=argparse.ArgumentParser()
    parser.add_argument('--local_rank',type=int, default=0)  #local_rank指定了輸出裝置，預設為GPU可用清單中的第一個GPU，必須加上。
    ...
    args = parser.parser.parse_args()
    return args

主函數寫：
def main():
    args = parse()
    torch.cuda.set_device(args.local_rank)  #必須寫在下一句的前面
   torch.distributed.init_process_group(
       'nccl',
       init_method='env://')

導入資料接口，需要用DistributedSampler
    dataset = ...
    num_workers = 4 if cuda else 0
    train_sampler=torch.utils.data.distributed.DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=batchsize, shuflle=False, num_worker=num_workers,pin_memory=cuda, drop_last=True, sampler=train_sampler)

定義模型：
net = XXXNet(using_amp=True)
net.train()
net= convert_syncbn_model(net)
device=torch.device('cuda:{}'.format(args.local_rank))
net=net.to(device)

定義優化器，損失函數，定義優化器一定要把模型搬運到GPU之上
apt = Adam([{'params':params_low_lr,'lr':4e-5},
    {'params':params_high_lr,'lr':1e-4}],weight_decay=settings.WEIGHT_DECAY)
crit = nn.BCELoss().to(device)

多GPU設定import torch.nn.parallel.DistributedDataParallel as DDP
net,opt = amp.initialize(net,opt,opt_level='o1')
net=DDP(net,delay_allreduce=True)loss使用方法：opt.zero_grad()with amp.scale_loss(loss, opt) as scaled_loss:    scaled_loss.backward()opt.step()加入主入口：if __name__ == '__main__':    main()無論是apex支援的DDP還是pytorch自身支援的DDP,都需使用torch.distributed.launch來使用，方法如下：CUDA_VISIBLE_DIVECES=1,2,4 python -m torch.distributed.launch --nproc_per_node=3 train.py1,2,4是GPU編号，nproc_per_node是指定用了哪些GPU,記得開頭說的local_rank，是因為torch.distributed.launch會調用這個local_ran

分布式訓練時儲存模型注意點：

如果直接在代碼中寫torch.save來儲存模型，則每個程序都會儲存一次相同的模型，會存在寫檔案寫到一半，會被個程序寫覆寫的情況。如何避免呢?

可以用local_rank == 0來僅僅在第一個GPU上執行程序來儲存模型檔案。

雖然是多個程序，但每個程序上模型的參數值都是一樣的，而預設代号為０的程序是主程序if arg.local_rank == 0:
    torch.save(xxx)

２．pytorch1.6及以上版本

有兩個接口：autocast和Gradscaler

  1. autocast

導入pytorch中子產品torch.cuda.amp的類autocast

from torch.cuda.amp import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

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

  with autocast():
    output=model(input)
    loss = loss_fn(output,target)

  loss.backward()
  optimizer.step()

可以使用autocast的context managers語義（如上），也可以使用decorators語義。當進入autocast上下文後，在這之後的cuda ops會把tensor的資料類型轉換為半精度浮點型，進而在不損失訓練精度的情況下加快運算。而不需要手動調用.half(),架構會自動完成轉換。

不過，autocast上下文隻能包含網絡的前向過程(包括loss的計算），不能包含反向傳播，因為BP的op會使用和前向op相同的類型。

當然，有時在autocast中的代碼會報錯：

Traceback (most recent call last):
......
 File "/opt/conda/lib/python3.8/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 scaler type float but found c10:Half,應該是個bug,可在tensor上手動調用.float()來讓type比對。

２）GradScaler

使用前，需要在訓練最開始前執行個體化一個GradScaler對象，例程如下：

from torch.cuda.amp import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

scaler = GradScaler() #訓練前執行個體化一個GradScaler對象

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

    with autocast():　＃前後開啟autocast
      output=model(input)
      loss = loss_fn(output,targt)

    scaler.scale(loss).backward()  #為了梯度放大
    #scaler.step()　首先把梯度值unscale回來，如果梯度值不是inf或NaN,則調用optimizer.step()來更新權重，否則，忽略step調用，進而保證權重不更新。　　 scaler.step(optimizer)
    scaler.update()  #準備着，看是否要增大scaler

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

１．如果出現inf或NaN,scaler.step(optimizer)會忽略此次權重更新(optimizer.step()），并将scaler的大小縮小（乘上backoff_factor)；

２．如果沒有出現inf或NaN,那麼權重正常更新，并且當連續多次(growth_interval指定)沒有出現inf或NaN，則scaler.update()會将scaler的大小增加(乘上growth_factor)。

對于分布式訓練，由于autocast是thread local的，要注意以下情形：

１）torch.nn.DataParallel：

以下代碼分布式是不生效的

model = MyModel()
dp_model = nn.DataParallel(model)

with autocast():
    output=dp_model(input)
loss=loss_fn(output)

需使用autocast裝飾model的forward函數

MyModel(nn.Module):
    @autocast()
    def forward(self, input):
        ...
        
#alternatively
MyModel(nn.Module):
    def forward(self, input):
        with autocast():
            ...


model = MyModel()
dp_model=nn.DataParallel(model)

with autocast():
    output=dp_model(input)
    loss = loss_fn(output)

2）torch.nn.DistributedDataParallel:

同樣，對于多GPU,也需要autocast裝飾model的forward方法，保證autocast在程序内部生效。

四. 注意事例：

在使用AMP時，由于報錯資訊并不明顯，給調試帶來了一定的難度。但隻要注意以下一些點，相信會少走很多彎路。

１．判斷GPU是否支援FP16，支援Tensor core的GPU（2080Ti,Titan,Tesla等），不支援的(Pascal系列）不建議；

1080Ti與2080Ti對比

gtx 1080ti:
半精度浮點數：0.17TFLOPS
單精度浮點數：11.34TFLOPS
雙精度浮點數：0.33TFLOPS
rtx 2080ti:
半精度浮點數：20.14TFLOPS
單精度浮點數：10.07TFLOPS
雙精度浮點數：0.31TFLOPS

半精度浮點數即FP16，單精度浮點數即FP32，雙精度浮點數即FP64。

在不使用apex的pytorch訓練過程中，一般預設均為單精度浮點數，從上面的資料可以看到1080ti和2080ti的單精度浮點數運算能力差不多，是以不使用apex時用1080ti和2080ti訓練模型時間上差别很小。

使用apex時用1個2080ti訓練時一個epoch是2h31min，兩者時間幾乎一樣，但是卻少用了一張2080ti。這是因為在pytorch訓練中使用apex時，此時大多數運算均為半精度浮點數運算，而2080ti的半精度浮點數運算能力是其單精度浮點數運算能力的兩倍

２．常數範圍：為了保證計算不溢出，首先保證人工設定的常數不溢出。如epsilon,INF等；

３．Dimension最好是8的倍數:次元是８的倍數，性能最好；

４．涉及sum的操作要小心，容易溢出，softmax操作，建議用官方API，并定義成layer寫在模型初始化裡；

５．模型書寫要規範：自定義的Layer寫在模型初始化函數裡，graph計算寫在forward裡；

６．一些不常用的函數，使用前要注冊：amp.register_float_function(torch,‘sogmoid’)

７．某些函數不支援FP16加速，建議不要用；

８．需要操作梯度的子產品必須在optimizer的step裡，不然AMP不能判斷grad是否為NaN

模式識别與深度學習 語音研究 Python Pytorch
上一篇: PyTorch的自動混合精度（AMP）
下一篇: Pytorch中的自動混合精度

繼續閱讀