淺談混合精度訓練imagenet

零、序

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

一、引言

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

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

訓練

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

訓練異常

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

訓練nan

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

二、混合精度訓練

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

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

  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不同的分布式訓練速度對比

環(huán)境配置如下：

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

GPU 8XV100 32G

cuda 10.2

pytorch 1.7.

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

多機啟動(dòng)

#!/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 \
...

單機啟動(dòng)

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鏈接。

實(shí)驗一、num workers對于速度的影響

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

num workers

BatchSize

FP16

epoch time

實(shí)驗二、OMP和MKL對于速度的影響

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

OMP & MKL

num workers

BatchSize

FP16

epoch time

實(shí)驗三、單機和多機啟動(dòng)速度差異

單機和多機啟動(dòng)，對于模型的前向基本是沒(méi)有影響的， 主要的差異是在loader開(kāi)始執行的速度，多機比起單機啟動(dòng)要快2倍-5倍左右的時(shí)間。

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

實(shí)驗均在ResNet50和imagenet下面進(jìn)行的，LR隨著(zhù)BS變換和線(xiàn)性增長(cháng)，公式如下

實(shí)驗結果

模型FP16+BNFP32實(shí)驗記錄

模型

數據集

batchsize（所有卡的總數）

優(yōu)化器

LearningRate

top1@acc

很明顯可以發(fā)現，單存使用FP16進(jìn)行訓練，但是沒(méi)有loss縮放的情況下，當BS和LR都增大的時(shí)候，訓練是無(wú)法進(jìn)行的，直接原因就是因為L(cháng)R過(guò)大，導致模型更新的時(shí)候數值范圍溢出了，同理loss也就直接為NAN了，我嘗試把LR調小后發(fā)現，模型是可以正常訓練的，只是精度略有所下降。

Apex混合精度實(shí)驗記錄

模型

MODE

數據集

batchsize（所有卡的總數）

優(yōu)化器

LearningRate

top1@acc

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

AMP實(shí)驗記錄

模型

MODE

數據集

batchsize（所有卡的總數）

優(yōu)化器

LearningRate

top1@acc

Time

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

2-bit訓練，ACTNN

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

附上一個(gè)比較合理的收斂情況

正常收斂情況

正常收斂情況2

五、結論

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

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

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

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

代碼在這里，自行取用。

本文僅做學(xué)術(shù)分享，如有侵權，請聯(lián)系刪文。