混合精度訓練介紹

Mixed-Precision Training是指在深度學習AI模型訓練過程中不同的層Layer采用不同的數據精度進行訓練， 最終使得訓練過程中的資源消耗(GPU顯存，GPU 算力)降低， 同時保證訓練可收斂，模型精度與高精度FP32的結果接近。

CNN ResNet 混合精度訓練

1. 導入torch.cuda.amp package
   由于CNN訓練要求大量算力， 因此一般混合精度需要使用 NVIDIA Automatic Mixed Precision (AMP)包， NVIDIA的AMP以及集成到了Pyorch, 因此直接調用 torch.cuda.amp APIs.

混合精度主要用到 Loss-Scaling (損失縮放) + Auto-cast (自動精度選擇/轉換)

# Mixed-Precision Training
from torch.cuda.amp.grad_scaler import GradScaler
from torch.cuda.amp.autocast_mode import autocast

# 實例化一個GradeScaler對象
scaler = GradScaler()

2. 對Training Loop進行修改， 修改2個地方

• 添加autocast(): autocast是一個Python context Manager, autocast 作用區(qū)域的代碼在運行的時候會跟據OP的類型，自動轉換為預定義好的低精度類型 （比如FP16)
  *注意: autocast一般作用的代碼區(qū)域為 Forward, Backward 階段需要指定autocast, 因此在Forward階段不同的layer (op)以及被設置了各自的精度模式， 在Backward階段，采用和Forward相同的精度進行計算。

• 添加GradeScalar： GradeScalar的目的是對權重的梯度矩陣值進行縮放(擴大化), 因為一般情況下的值非常小，如果采用低精度類型 （FP16）,則導致下溢underflow. 解決方法之一就是希望將 進行縮放變大； 由于Loss函數導數具有線性性質， 因此也可以對 Loss進行縮放，實際等價于對梯度值進行了放大。

 for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # move data to the same device as model
        images = images.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        # compute output
        with autocast(enabled=args.mixed_precision, dtype=torch.float16):
                output = model(images)
                loss = criterion(output, target)

        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        # losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        # loss.backward()
        # optimizer.step()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
         ...

為了便于觀察訓練過程， 在代碼中添加了Pytorch Profiler進行可視化:
為了說明情況， 可以只跑少數的幾個batch即可

image.png

image.png

image.png

對應的CUDA kernel函數， FP16類型的

image.png

采用BFLOAT16進行混合精度

方法很簡單，在autocast的dtype設置為torch.bfloat16。 除此之外，需要采用支持BFLOAT16類型的計算設備（TPU， >=NVIDIA Ampere/Volta 架構的GPU, 比如NVIDIA V100, A100, RTX 30/40系列）

with autocast(enabled=args.mixed_precision, dtype=torch.bfloat16):
                output = model(images)
                loss = criterion(output, target)

image.png

image.png