pytorch设置apex混合精度 pytorch amp自动混合精度

转载

mob6454cc777577 2024-01-15 07:37:50

文章标签 pytorch设置apex混合精度深度学习 pytorch python 权重 文章分类 PyTorch 人工智能

AMP介绍

之前大多数的学习框架都是用float32的精度进行训练，AMP做的改进就是使用float32和float16相结合进行训练，并且相同的超参数实现了与使用float32几乎相同的精度

为什么使用混合精度

混合精度预示着有不止一种精度的Tensor，PyTorch的AMP有2种精度是torch.FloatTensor和torch.HalfTensor

因为使用自动混合精度其实一种在训练过程中对训练速度和显存的优化，众所周知神经网络模型的训练通常会大量占据显存，训练时间也特别长，因此做出优化是很有必要的。

调用方法是torch.cuda.amp.autocast()和torch.cuda.amp.GradScaler() 字面意思看只能在cuda上使用，事实上，这个功能正是NVIDIA的开发人员贡献到PyTorch项目中的。后面会用代码详细介绍这两种方法的使用。

代码

scaler = torch.cuda.amp.GradScaler()

def train_one_epoch(loader, model, loss_fn, optimizer, scaler):
    for batch_idx, (data, targets) in enumerate(loop):
     /
     .......
     /
     // scaler的大小在每次迭代中动态估计，为了尽可能减少梯度，scaler应该更大；但太大，半精度浮点型又容易overflow（变成inf或NaN).所以，动态估计原理就是在不出现if或NaN梯度的情况下，尽可能的增大scaler值。在每次scaler.step(optimizer)中，都会检查是否有inf或NaN的梯度出现：
        with torch.cuda.amp.autocast():
            scores = model(data)
            loss = loss_fn(scores, targets)
        optimizer.zero_grad()
        scaler.scale(loss).backward() //反向传播梯度放大
        scaler.step(optimizer)//首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。
        scaler.update()//看是否要增大scaler

torch.HalfTensor

torch.HalfTensor的优势就是存储小、计算快、更好的利用CUDA设备的Tensor Core。因此训练的时候可以减少显存的占用，同时训练速度更快。

torch.HalfTensor的劣势就是溢出错误，在开启 autocast 后，网络中间结果的数值类型会变成 float16，其对应的梯度自然也是 float16。而由于 float16 的数值范围比 float32 要小，所以如果遇到特别小的数（比如 loss、gradient），float16 就难以精确表达，这时候一般需要进行梯度放缩（Gradient Scaling）。做法是对网络的 loss 进行放大，从而使反向传播时网络中间结果对应的梯度也得到相同的放大，减少精度的损失，而梯度反传到参数时，仍会是 float32 的类型，在等比缩小之后，并不会影响参数的更新。