Pytorch自动混合精度(AMP)的使用总结

Posted AI浩

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了Pytorch自动混合精度(AMP)的使用总结相关的知识,希望对你有一定的参考价值。

Pytorch自动混合精度(AMP)的使用总结

pytorch从1.6版本开始,已经内置了torch.cuda.amp,采用自动混合精度训练就不需要加载第三方NVIDIA的apex库了。本文借鉴别人的文章和自己的经验编写,如果有错误还请大家指正。

本文包含一下内容:

1、介绍混合精度实现的两个接口。

2、如何将混合精度和梯度裁剪结合。

3、如果在torch.nn.DataParallel方式下实现混合精度训练。

以上几个问题,我也被困扰了好久,写这篇文章记录一下。

pytorch实现混合精度的两个接口

pytorch实现混合精度有两个接口:autocast和Gradscaler。

1、使用autocast,导入pytorch中模块torch.cuda.amp的类autocast

from torch.cuda.amp import autocast as autocast

model=mobilenetv2().cuda()
optimizer=optim.SGD(model.parameters(),...)

for input,target in data:
  optimizer.zero_grad()
  with autocast():
    output=model(input)
    loss = loss_fn(output,target)
  loss.backward()
  optimizer.step()

当进入autocast上下文后,在这之后的cuda ops会把tensor的数据类型转换为半精度浮点型,从而在不损失训练精度的情况下加快运算。而不需要手动调用.half(),框架会自动完成转换。不过,autocast上下文只能包含网络的前向过程(包括loss的计算),不能包含反向传播,因为BP的op会使用和前向op相同的类型。

有时在autocast中的代码会报错,错误如下:

Traceback (most recent call last):
......
 File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
  result = self.forward(*input, ** kwargs)
......
RuntimeError: expected scalar type float but found c10::Half

对于RuntimeError:expected scaler type float but found c10:Half,应该是个bug,可在tensor上手动调用.float()来让type匹配。

2、GradScaler,使用前,需要在训练最开始前实例化一个GradScaler对象,例程如下:

from torch.cuda.amp import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

scaler = GradScaler() #训练前实例化一个GradScaler对象

for epoch in epochs:
  for input,target in data:
    optimizer.zero_grad()

    with autocast(): #前后开启autocast
      output=model(input)
      loss = loss_fn(output,targt)

    scaler.scale(loss).backward()  #为了梯度放大
    #scaler.step() 首先把梯度值unscale回来,如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重,否则,忽略step调用,从而保证权重不更新。   
    scaler.step(optimizer)
    scaler.update()  #准备着,看是否要增大scaler

scaler的大小在每次迭代中动态估计,为了尽可能减少梯度underflow,scaler应该更大;但太大,半精度浮点型又容易overflow(变成inf或NaN).所以,动态估计原理就是在不出现if或NaN梯度的情况下,尽可能的增大scaler值。在每次scaler.step(optimizer)中,都会检查是否有inf或NaN的梯度出现:

1.如果出现inf或NaN,scaler.step(optimizer)会忽略此次权重更新(optimizer.step()),并将scaler的大小缩小(乘上backoff_factor);

2.如果没有出现inf或NaN,那么权重正常更新,并且当连续多次(growth_interval指定)没有出现inf或NaN,则scaler.update()会将scaler的大小增加(乘上growth_factor)。

建议使用第二种方式,我在写模型时,使用的第二种方式。

混合精度与梯度裁剪

如何在混合精度下实现梯度。

from torch.cuda.amp import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

scaler = GradScaler() #训练前实例化一个GradScaler对象
for epoch in epochs:
  for input,target in data:
    optimizer.zero_grad()

    with autocast(): #前后开启autocast
      output=model(input)
      loss = loss_fn(output,targt)
    scaler.scale(loss).backward()  #为了梯度放大
    scaler.unscale_(optimizer)
	torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)  
    scaler.step(optimizer)
    scaler.update()  #准备着,看是否要增大scaler

unscale_()是可选的,用于需要修改或检查后向传递和 step() 之间的梯度的情况。如果 unscale_() 没有被显式调用,渐变将在 [step() 期间自动取消缩放。

unscale_() 只能在每个优化器每次 step() 调用时调用一次,并且只有在该优化器分配参数的所有梯度都已累积之后。在每个 step() 之间为给定优化器调用 unscale_()两次会触发 RuntimeError。

torch.nn.DataParallel方式下实现混合精度训练

默认的情况下,混合精度在torch.nn.DataParallel模型是不生效的,需要用需使用autocast修饰model的forward函数。例:

from torch.cuda.amp import autocast

MyModel(nn.Module):
    @autocast()
    def forward(self, input):
        ...
        
#或者
MyModel(nn.Module):
    def forward(self, input):
        with autocast():
            ...

model = MyModel()
dp_model=nn.DataParallel(model)

with autocast():
    output=dp_model(input)
    loss = loss_fn(output)

这样就可以实现多显卡混合精度训练了。不过如果不适用混合精度的时候,要把这个模型的autocast注释掉。否则,loss一直是nan。

参考:

Pytorch自动混合精度(AMP)介绍与使用 - jimchen1218 - 博客园 (cnblogs.com)

以上是关于Pytorch自动混合精度(AMP)的使用总结的主要内容,如果未能解决你的问题,请参考以下文章

Pytorch AMP——自动混合精度训练

PyTorch 半精度训练踩坑

实现混合精度的两种方式

NLP涉及技术原理和应用简单讲解:paddle(分布式训练AMP自动混合精度训练模型量化模型性能分析)

NLP涉及技术原理和应用简单讲解:paddle(分布式训练AMP自动混合精度训练模型量化模型性能分析)

混合,带有多个标签的精度