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

举报
AI浩 发表于 2022/08/17 09:41:35 2022/08/17
【摘要】 Pytorch自动混合精度(AMP)的使用总结pytorch从1.6版本开始,已经内置了torch.cuda.amp,采用自动混合精度训练就不需要加载第三方NVIDIA的apex库了。本文借鉴别人的文章和自己的经验编写,如果有错误还请大家指正。本文包含一下内容:1、介绍混合精度实现的两个接口。2、如何将混合精度和梯度裁剪结合。3、如果在torch.nn.DataParallel方式下实现混...

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)

【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息, 否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。