淺談混合精度訓(xùn)練imagenet

零、序

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

一、引言

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

這種是比較正常的損失和acc的情況，因?yàn)轫?xiàng)目的數(shù)據(jù)非常長(zhǎng)尾。

訓(xùn)練

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

訓(xùn)練異常

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

訓(xùn)練nan

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

二、混合精度訓(xùn)練

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

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

  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庫(kù)，這里有O1,O2,O3三種訓(xùn)練模式，代碼如下：

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版本以后把a(bǔ)pex并入到了自身的庫(kù)里面，代碼如下：

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不同的分布式訓(xùn)練速度對(duì)比

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

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

GPU 8XV100 32G

cuda 10.2

pytorch 1.7.

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

多機(jī)啟動(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 \
...

單機(jī)啟動(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 \
...

詳細(xì)代碼看文末的github鏈接。

實(shí)驗(yàn)一、num workers對(duì)于速度的影響

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

num workers

BatchSize

FP16

epoch time

實(shí)驗(yàn)二、OMP和MKL對(duì)于速度的影響

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

OMP & MKL

num workers

BatchSize

FP16

epoch time

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

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

四、不同混合精度訓(xùn)練方法對(duì)比

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

實(shí)驗(yàn)結(jié)果

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

模型

數(shù)據(jù)集

batchsize（所有卡的總數(shù)）

優(yōu)化器

LearningRate

top1@acc

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

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

模型

MODE

數(shù)據(jù)集

batchsize（所有卡的總數(shù)）

優(yōu)化器

LearningRate

top1@acc

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

AMP實(shí)驗(yàn)記錄

模型

MODE

數(shù)據(jù)集

batchsize（所有卡的總數(shù)）

優(yōu)化器

LearningRate

top1@acc

Time

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

2-bit訓(xùn)練，ACTNN

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

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

正常收斂情況

正常收斂情況2

五、結(jié)論

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

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

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

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

代碼在這里，自行取用。

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