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的梯度出现：

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

２．如果没有出现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)