AI编辑:我是小将
混合精度训练(mixed precision training)可以让模型训练在尽量不降低性能的情形下提升训练速度,而且也可以降低显卡使用内存。目前主流的深度学习框架都开始支持混合精度训练。对于PyTorch,混合精度训练还主要是采用NVIDIA开源的apex库。但是,PyTorch将迎来重大更新,那就是提供内部支持的混合精度训练,而且是自动混合精度训练:
torch.cuda.amp.autocast
:自动为GPU op选择精度来提升训练性能而不降低模型准确度。
torch.cuda.amp.GradScaler
:对梯度进行scale来加快模型收敛,因为float16梯度容易出现underflow(梯度过小)
两者结合在一起,可以实现自动混合精度训练:
model=Net().cuda() optimizer=optim.SGD(model.parameters(),...) #CreatesaGradScaleronceatthebeginningoftraining. scaler=GradScaler() forepochinepochs: forinput,targetindata: optimizer.zero_grad() #Runstheforwardpasswithautocasting. withautocast(): output=model(input) loss=loss_fn(output,target) #Scalesloss.Callsbackward()onscaledlosstocreatescaledgradients. #Backwardpassesunderautocastarenotrecommended. #Backwardopsruninthesameprecisionthatautocastusedforcorrespondingforwardops. scaler.scale(loss).backward() #scaler.step()firstunscalesthegradientsoftheoptimizer'sassignedparams. #IfthesegradientsdonotcontaininfsorNaNs,optimizer.step()isthencalled, #otherwise,optimizer.step()isskipped. scaler.step(optimizer) #Updatesthescalefornextiteration. scaler.update()#Createsmodelandoptimizerindefaultprecision
可以看到,为了防止梯度的underflow,首先scaler.scale(loss).backward()会对loss乘以一个scale因子,然后backward时所有梯度都会乘以相同的scale因子,这样保证梯度有较大的magnitude而不会出现为0。我们不希望这个scale因子对学习速率产生影响,那么scaler.step(optimizer)会先unscale要更新的梯度然后再更新,如果梯度出现infs或者NaNs,optimizer将忽略这次迭代训练。
如果你想在梯度更新前对梯度进行clip,也是可以的:
forepochinepochs: forinput,targetindata: optimizer.zero_grad() withautocast(): output=model(input) loss=loss_fn(output,target) scaler.scale(loss).backward() #Unscalesthegradientsofoptimizer'sassignedparamsin-place scaler.unscale_(optimizer) #Sincethegradientsofoptimizer'sassignedparamsareunscaled,clipsasusual: torch.nn.utils.clip_grad_norm_(model.parameters(),max_norm) #optimizer'sgradientsarealreadyunscaled,soscaler.stepdoesnotunscalethem, #althoughitstillskipsoptimizer.step()ifthegradientscontaininfsorNaNs. scaler.step(optimizer) #Updatesthescalefornextiteration. scaler.update()scaler=GradScaler()
当然,混合精度训练肯定要支持分布式训练,由于autocast
是thread local的,所以要注意以下不同的情形:
如果使用torch.nn.DataParallel
:
此时只有一个进程,而不同GPU上是各自的线程跑forward过程的,所以下面操作时无效的:
dp_model=nn.DataParallel(model) #Setsautocastinthemainthread withautocast(): #dp_model'sinternalthreadswon'tautocast.Themainthread'sautocaststatehasnoeffect. output=dp_model(input) #loss_fnstillautocasts,butit'stoolate... loss=loss_fn(output)model=MyModel()
此时你需要对model的forward方法用autocast
装饰:
... @autocast() defforward(self,input): ... #Alternatively MyModel(nn.Module): ... defforward(self,input): withautocast(): ... model=MyModel() dp_model=nn.DataParallel(model) withautocast(): output=dp_model(input) loss=loss_fn(output)MyModel(nn.Module):
如果使用torch.nn.parallel.DistributedDataParallel
:
一般情形下是单GPU进程的,此时原来的用来就没有问题,但是如果是多GPU一个进程那么就和上述问题一样,需要用autocast装饰model的forward。
更多内容见:/docs/master/notes/amp_examples.html#amp-examples
推荐阅读
堪比Focal Loss!解决目标检测中样本不平衡的无采样方法
超越BN和GN!谷歌提出新的归一化层:FRN
CVPR 论文系列一:VSGNet
PyTorch分布式训练简明教程
图神经网络入门
另辟蹊径!斯坦福大学提出边界框回归任务新Loss:GIoU
人人必须要知道的语义分割模型:DeepLabv3+
机器学习算法工程师
一个用心的公众号