混合精度訓(xùn)練介紹
Mixed-Precision Training是指在深度學(xué)習(xí)AI模型訓(xùn)練過程中不同的層Layer采用不同的數(shù)據(jù)精度進(jìn)行訓(xùn)練仿村， 最終使得訓(xùn)練過程中的資源消耗(GPU顯存绕娘，GPU 算力)降低缴啡， 同時(shí)保證訓(xùn)練可收斂纠修，模型精度與高精度FP32的結(jié)果接近撵溃。
CNN ResNet 混合精度訓(xùn)練
-
導(dǎo)入torch.cuda.amp package
由于CNN訓(xùn)練要求大量算力绣否， 因此一般混合精度需要使用 NVIDIA Automatic Mixed Precision (AMP)包泞当， NVIDIA的AMP以及集成到了Pyorch, 因此直接調(diào)用torch.cuda.ampAPIs.
混合精度主要用到 Loss-Scaling (損失縮放) + Auto-cast (自動(dòng)精度選擇/轉(zhuǎn)換)
# Mixed-Precision Training
from torch.cuda.amp.grad_scaler import GradScaler
from torch.cuda.amp.autocast_mode import autocast
# 實(shí)例化一個(gè)GradeScaler對(duì)象
scaler = GradScaler()
- 對(duì)Training Loop進(jìn)行修改， 修改2個(gè)地方
添加
autocast(): autocast是一個(gè)Python context Manager, autocast 作用區(qū)域的代碼在運(yùn)行的時(shí)候會(huì)跟據(jù)OP的類型奇颠，自動(dòng)轉(zhuǎn)換為預(yù)定義好的低精度類型 （比如FP16)
*注意: autocast一般作用的代碼區(qū)域?yàn)?Forward, Backward 階段需要指定autocast, 因此在Forward階段不同的layer (op)以及被設(shè)置了各自的精度模式败去， 在Backward階段，采用和Forward相同的精度進(jìn)行計(jì)算烈拒。添加
GradeScalar： GradeScalar的目的是對(duì)權(quán)重的梯度矩陣值進(jìn)行縮放(擴(kuò)大化), 因?yàn)橐话闱闆r下
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()
...
為了便于觀察訓(xùn)練過程伴郁， 在代碼中添加了Pytorch Profiler進(jìn)行可視化:
為了說明情況耿战， 可以只跑少數(shù)的幾個(gè)batch即可
對(duì)應(yīng)的CUDA kernel函數(shù)， FP16類型的
采用BFLOAT16進(jìn)行混合精度
方法很簡單焊傅，在autocast的dtype設(shè)置為torch.bfloat16剂陡。 除此之外，需要采用支持BFLOAT16類型的計(jì)算設(shè)備（TPU狐胎， >=NVIDIA Ampere/Volta 架構(gòu)的GPU, 比如NVIDIA V100, A100, RTX 30/40系列）
with autocast(enabled=args.mixed_precision, dtype=torch.bfloat16):
output = model(images)
loss = criterion(output, target)
