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)的使用总结的主要内容,如果未能解决你的问题,请参考以下文章
NLP涉及技术原理和应用简单讲解:paddle(分布式训练AMP自动混合精度训练模型量化模型性能分析)