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。
参考：