pytorch 如何使用amp進行混合精度訓練

更新時間：2021年05月24日 11:59:15 作者：多財多億,憑億近人

這篇文章主要介紹了pytorch 使用amp進行混合精度訓練的操作，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教

簡介

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

在Pytorch 1.5版本及以前，通過NVIDIA提供的apex庫可以實現(xiàn)amp功能。但是在使用過程中會伴隨著一些版本兼容和奇怪的報錯問題。

從1.6版本開始，Pytorch原生支持自動混合精度訓練，并已進入穩(wěn)定階段，AMP 訓練能在 Tensor Core GPU 上實現(xiàn)更高的性能并節(jié)省多達 50％的內(nèi)存。

環(huán)境

Python 3.8

Pytorch 1.7.1

CUDA 11 + cudnn 8

NVIDIA GeFore RTX 3070

ps：后續(xù)使用移動端的3070，或者3080結合我目前訓練的分類網(wǎng)絡來測試實際效果

原理

關于低精度計算

當前的深度學習框架大都采用的都是FP32來進行權重參數(shù)的存儲，比如Python float的類型為雙精度浮點數(shù) FP64，PyTorch Tensor的默認類型為單精度浮點數(shù)FP32。

隨著模型越來越大，加速訓練模型的需求就產(chǎn)生了。在深度學習模型中使用FP32主要存在幾個問題，第一模型尺寸大，訓練的時候對顯卡的顯存要求高；第二模型訓練速度慢；第三模型推理速度慢。

其解決方案就是使用低精度計算對模型進行優(yōu)化。

推理過程中的模型優(yōu)化目前比較成熟的方案就是FP16量化和INT8量化，NVIDIA TensorRT等框架就可以支持，這里不再贅述。訓練方面的方案就是混合精度訓練，它的基本思想很簡單: 精度減半(FP32→ FP16) ，訓練時間減半。

與單精度浮點數(shù)float32（32bit，4個字節(jié)）相比，半精度浮點數(shù)float16僅有16bit，2個字節(jié)組成。

可以很明顯的看到，使用FP16可以解決或者緩解上面FP32的兩個問題：顯存占用更少：通用的模型FP16占用的內(nèi)存只需原來的一半，訓練的時候可以使用更大的batchsize。

計算速度更快：有論文指出半精度的計算吞吐量可以是單精度的 2-8 倍。

從上到下依次為 fp16、fp32 、fp64

當前很多NVIDIA GPU搭載了專門為快速FP16矩陣運算設計的特殊用途Tensor Core，比如Tesla P100，Tesla V100、Tesla A100、GTX 20XX 和RTX 30XX等。

Tensor Core是一種矩陣乘累加的計算單元，每個Tensor Core每個時鐘執(zhí)行64個浮點混合精度操作（FP16矩陣相乘和FP32累加），英偉達宣稱使用Tensor Core進行矩陣運算可以輕易的提速，同時降低一半的顯存訪問和存儲。

隨著Tensor Core的普及FP16計算也一步步走向成熟，低精度計算也是未來深度學習的一個重要趨勢。

Tensor Core 的 4x4x4 矩陣乘法與累加

Volta GV100 Tensor Core 流程圖

自動混合精度訓練

不同于在推理過程中直接削減權重精度，在模型訓練的過程中，直接使用半精度進行計算會導致的兩個問題的處理：舍入誤差(Rounding Error)和溢出錯誤(Grad Overflow / Underflow)。

舍入誤差： float16的最大舍入誤差約為 (~2 ^-10 )，比float32的最大舍入誤差(~2 ^-23) 要大不少。對足夠小的浮點數(shù)執(zhí)行的任何操作都會將該值四舍五入到零，在反向傳播中很多甚至大多數(shù)梯度更新值都非常小，但不為零。在反向傳播中舍入誤差累積可以把這些數(shù)字變成0或者 nan，這會導致不準確的梯度更新，影響網(wǎng)絡的收斂。

溢出錯誤： 由于float16的有效的動態(tài)范圍約為 ( 5.96×10^-8 ~ 6.55×10^4)，比單精度的float32(1.4x10^-45 ~ 1.7x10^38)要狹窄很多，精度下降(小數(shù)點后16相比較小數(shù)點后8位要精確的多)會導致得到的值大于或者小于fp16的有效動態(tài)范圍，也就是上溢出或者下溢出。

在深度學習中，由于激活函數(shù)的的梯度往往要比權重梯度小，更易出現(xiàn)下溢出的情況。2018年ICLR論文 Mixed Precision Training 中提到，簡單的在每個地方使用FP16會損失掉梯度更新小于2^-24的值——大約占他們的示例網(wǎng)絡所有梯度更新的5%。

解決方案就是使用混合精度訓練（Mixed Precision）和損失縮放(Loss Scaling)：

1、混合精度訓練：

混合精度訓練是一種通過在FP16上執(zhí)行盡可能多的操作來大幅度減少神經(jīng)網(wǎng)絡訓練時間的技術，在像線性層或是卷積操作上，F(xiàn)P16運算較快，但像Reduction運算又需要 FP32的動態(tài)范圍。通過混合精度訓練的方式，便可以在部分運算操作使用FP16，另一部分則使用 FP32，混合精度功能會嘗試為每個運算使用相匹配的數(shù)據(jù)類型，在內(nèi)存中用FP16做儲存和乘法從而加速計算，用FP32做累加避免舍入誤差。這樣在權重更新的時候就不會出現(xiàn)舍入誤差導致更新失敗，混合精度訓練的策略有效地緩解了舍入誤差的問題。

2、損失縮放：

即使用了混合精度訓練，還是會存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出。損失縮放是指在執(zhí)行反向傳播之前，將損失函數(shù)的輸出乘以某個標量數(shù)(論文建議從8開始)。乘性增加的損失值產(chǎn)生乘性增加的梯度更新值，提升許多梯度更新值到超過FP16的安全閾值2^-24。只要確保在應用梯度更新之前撤消縮放，并且不要選擇一個太大的縮放以至于產(chǎn)生inf權重更新(上溢出) ，從而導致網(wǎng)絡向相反的方向發(fā)散。

使用Pytorch AMP

Pytorch原生的amp模式使用起來相當簡單，只需要從torch.cuda.amp導入GradScaler和 autocast這兩個函數(shù)即可。torch.cuda.amp的名字意味著這個功能只能在cuda上使用，事實上，這個功能正是NVIDIA的開發(fā)人員貢獻到PyTorch項目中的。

Pytorch在amp模式下維護兩個權重矩陣的副本，一個主副本用 FP32，一個半精度副本用 FP16。梯度更新使用FP16矩陣計算，但更新于 FP32矩陣。這使得應用梯度更新更加安全。

autocast上下文管理器實現(xiàn)了 FP32到FP16的轉換，它會自動判別哪些層可以進行FP16哪些層不可以。 GradScaler對梯度更新計算(檢查是否溢出)和優(yōu)化器(將丟棄的batches轉換為 no-op)進行控制，通過放大loss的值來防止梯度的溢出。

在訓練中的具體使用方法如下所示：

def train():
    batch_size = 8
    epochs = 10
    lr = 1e-3
    size = 256
    num_class = 35
    use_amp = True 
    device = 'cuda' if torch.cuda.is_available() else 'cpu' 
    print('torch version: {}'.format(torch.__version__))
    print('amp:           {}'.format(use_amp))
    print('device:        {}'.format(device))
    print('epochs:        {}'.format(epochs))
    print('learn rate:    {}'.format(lr))
    print('batch size:    {}'.format(batch_size))
 
    net = ERFNet(num_classes=num_class).to(device) 
    train_data = CityScapesDataset('D:\\dataset\\cityscapes',
                                   'D:\\dataset\\cityscapes\\trainImages.txt',
                                   'D:\\dataset\\cityscapes\\trainLabels.txt',
                                   size, num_class)
    val_data = CityScapesDataset('D:\\dataset\\cityscapes',
                                 'D:\\dataset\\cityscapes\\valImages.txt',
                                 'D:\\dataset\\cityscapes\\valLabels.txt',
                                 size, num_class)
 
    train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=False, num_workers=8)
    val_dataloader = DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=4)
 
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    criterion = torch.nn.CrossEntropyLoss(ignore_index=255)
 
    if use_amp:
        scaler = torch.cuda.amp.GradScaler() 
    writer = SummaryWriter("summary") 
    train_loss = AverageMeter()
    val_acc = AverageMeter()
    val_miou = AverageMeter()
 
    for epoch in range(0, epochs):
        train_loss.reset()
        val_acc.reset()
        val_miou.reset()
 
        with tqdm(total=train_data.__len__(), unit='img', desc="Epoch {}/{}".format(epoch + 1, epochs)) as pbar:
            # train
            net.train()
            for img, mask in train_dataloader:
                img = img.to(device)
                mask = mask.to(device)
                n = img.size()[0] 
                opt.zero_grad()
 
                if use_amp:
                    with torch.cuda.amp.autocast():
                        output = net(img)
                        loss = criterion(output, mask)
 
                    scaler.scale(loss).backward()
                    scaler.step(opt)
                    scaler.update()
                else:
                    output = net(img)
                    loss = criterion(output, mask)
                    loss.backward()
                    opt.step() 
                train_loss.update(loss.item(), n)
 
                pbar.set_postfix(**{"loss": train_loss.avg})
                pbar.update(img.size()[0])
 
            writer.add_scalar('train_loss', train_loss.avg, epoch)
            # eval
            net.eval()
            for img, mask in val_dataloader:
                img = img.to(device)
                mask = mask
                n = img.size()[0] 
                output = net(img) 
                pred_mask = torch.softmax(output, dim=1)
                pred_mask = pred_mask.detach().cpu().numpy()
                pred_mask = np.argmax(pred_mask, axis=1)
                true_mask = mask.numpy()
                acc, acc_cls, mean_iu, fwavacc = evaluate(pred_mask, true_mask, num_class)
 
                val_acc.update(acc)
                val_miou.update(mean_iu)
 
            writer.add_scalar('val_acc', val_acc.avg, epoch)
            writer.add_scalar('val_miou', val_miou.avg, epoch) 
            pbar.set_postfix(**{"loss": train_loss.avg, "val_acc": val_acc.avg, "val_miou": val_miou.avg})