淺談混合精度訓練imagenet

零、序

本文沒有任何的原理和解讀，只有一些實驗的結論，對于想使用混合精度訓練的同學可以直接參考結論白嫖，或者直接拿github上的代碼(文末放送)。

一、引言

以前做項目的時候出現過一個問題，使用FP16訓練的時候，只要BatchSize增加(LR也對應增加)的時候訓練，一段時間后就會出現loss異常，同時val對應的明顯降低，甚至直接NAN的情況出現，圖示如下：

這種是比較正常的損失和acc的情況，因為項目的數據非常長尾。

訓練

這種就是不正常的訓練情況, val的損失不下降反而上升，acc不升反而降。

訓練異常

還有一種情況，就是訓練十幾個epoch以后，loss上升到非常大，acc為nan，后續訓練都是nan，tensorboard顯示有點問題，只好看ckpt的結果了。

訓練nan

由于以前每周都沒跑很多模型，問題也不是經常出現，所以以為是偶然時間，不過最近恰好最近要做一些transformer的實驗，在跑imagenet baseline(R50)的時候，出現了類似的問題，由于FP16訓練的時候，出現了溢出的情況所導致的。簡單的做了一些實驗，整理如下。

二、混合精度訓練

混合精度訓練，以pytorch 1.6版本為基礎的話，大致是有3種方案，依次介紹如下：

模型和輸入輸出直接half，如果有BN，那么BN計算需要轉為FP32精度，我上面的問題就是基于此來訓練的，代碼如下：

  if args.FP16:
        model = model.half()
        for bn in get_bn_modules(model):
            bn.float()
    ...
    for data in dataloader:
        if args.FP16:
            image, label = data[0].half()
            output = model(image)
            losses = criterion(output, label)
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

使用NVIDIA的Apex庫，這里有O1,O2,O3三種訓練模式，代碼如下：

try:
    from apex import amp 
    from apex.parallel import convert_syncbn_model
    from apex.parallel import DistributedDataParallel as DDP 
except Exception as e:
    print("amp have not been import !!!")
if args.apex:
   model = convert_syncbn_model(model)
if args.apex:
   model, optimizer = amp.initialize(model, optimizer, opt_level=args.mode) 
   model = DDP(model, delay_allreduce=True)
...
for data in dataloader:
    image, label = data[0], data[1]
    batch_output = model(image)
    losses = criterion(batch_output, label)
    optimizer.zero_grad()
    if args.apex:
        with amp.scale_loss(losses, optimizer) as scaled_loss:
            scaled_loss.backward()
        optimizer.step()

pytorch1.6版本以后把apex并入到了自身的庫里面，代碼如下：

from torch.cuda.amp import autocast as autocast
from torch.nn.parallel import DistributedDataParallel as DataParallel
model = DataParallel(model, 
                        device_ids=[args.local_rank], 
                        find_unused_parameters=True)
if args.amp:
        scaler = torch.cuda.amp.GradScaler()
for data in dataloader:
    image, label = data[0], data[1]
    if args.amp:
        with autocast():
            batch_output = model(image)
            losses = criterion(batch_output, label)
    if args.amp:
        scaler.scale(losses).backward()
        scaler.step(optimizer)
        scaler.update()

三、pytorch不同的分布式訓練速度對比

環境配置如下：

CPU Intel(R) Xeon(R) Platinum 8163 CPU @ 2.50GHz

GPU 8XV100 32G

cuda 10.2

pytorch 1.7.

pytorch分布式有兩種不同的啟動方法，一種是單機多卡啟動，一種是多機多卡啟動, ps: DataParallel不是分布式訓練。

多機啟動

#!/bin/bash
cd $FOLDER;
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -W ignore -m torch.distributed.launch --nproc_per_node 8 train_lanuch.py \
...

單機啟動

cd $FOLDER;

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -W ignore test.py \
--dist-url 'tcp://127.0.0.1:9966' \
--dist-backend 'nccl' \
--multiprocessing-distributed=1 \
--world-size=1 \
--rank=0 \
...

詳細代碼看文末的github鏈接。

實驗一、num workers對于速度的影響

我的服務器是48個物理核心，96個邏輯核心，所以48的情況下，效果最好，不過增加和減少對于模型的影響不大，基本上按照CPU的物理核心個數來設置就可以。

num workers

BatchSize

FP16

epoch time

實驗二、OMP和MKL對于速度的影響

OMP和MKL對于多機模式下的速度有輕微的影響，如果不想每個都去試，直接經驗設置為1最合理。FP16大幅度提升模型的訓練速度，可以節省2/5的時間。

OMP & MKL

num workers

BatchSize

FP16

epoch time

實驗三、單機和多機啟動速度差異

單機和多機啟動，對于模型的前向基本是沒有影響的， 主要的差異是在loader開始執行的速度，多機比起單機啟動要快2倍-5倍左右的時間。

四、不同混合精度訓練方法對比

實驗均在ResNet50和imagenet下面進行的，LR隨著BS變換和線性增長，公式如下

實驗結果

模型FP16+BNFP32實驗記錄

模型

數據集

batchsize（所有卡的總數）

優化器

LearningRate

top1@acc

很明顯可以發現，單存使用FP16進行訓練，但是沒有loss縮放的情況下，當BS和LR都增大的時候，訓練是無法進行的，直接原因就是因為LR過大，導致模型更新的時候數值范圍溢出了，同理loss也就直接為NAN了，我嘗試把LR調小后發現，模型是可以正常訓練的，只是精度略有所下降。

Apex混合精度實驗記錄

模型

MODE

數據集

batchsize（所有卡的總數）

優化器

LearningRate

top1@acc

Apex O3模式下的訓練情況和上面FP16的結論是一致的，存FP16訓練，不管是否有loss縮放都會導致訓練NaN，O2和O1是沒有任何問題的，O2的精度略低于O1的精度。

AMP實驗記錄

模型

MODE

數據集

batchsize（所有卡的總數）

優化器

LearningRate

top1@acc

Time

AMP自動把模型需要用FP32計算的層或者op直接轉換，不需要顯著性指定。精度比apex高，同時訓練時間更少。

2-bit訓練，ACTNN

簡單的嘗試了一下2bit訓練，1k的bs是可以跑的，不過速度相比FP16跑，慢了太多，基本可以pass掉了。

附上一個比較合理的收斂情況

正常收斂情況

正常收斂情況2

五、結論

如果使用分布式訓練，使用pytorch 多機模式啟動，收益比較高，如果你不希望所有卡都用的話，那么建議使用單機多卡的模式。

如果使用FP16方式計算的話，那么無腦pytorch amp就可以了，速度和精度都比較有優勢，代碼量也不多。

我的增強只用了隨機裁剪，水平翻轉，跑了90個epoch，原版的resnet50是跑了120個epoch，還有color jitter，imagenet上one crop的結果0.76012，和我的結果相差無幾，所以分類任務(基本上最后是求概率的問題，圖像，視頻都work，已經驗證過)上FP16很明顯完全可以替代FP32。我跑了一個120epoch的版本，結果是0.767，吊打原版本結果了QAQ。

如果跑小的bs，第一種FP16的方法完全是ok的，對于大的bs來說，使用AMP會使得模型的收斂更加穩定。

代碼在這里，自行取用。

本文僅做學術分享，如有侵權，請聯系刪文。