NVIDIA-apex 和 Pytoch-amp学习和理解

一、apex是什么

1、apex和amp

英伟达官网https://nvidia.github.io/apex/index.html可以看到apex的全称：A PyTorch Extension(Apex)，其实就是一种pytorch的拓展插件。

目的就是使用户能够快速实现amp（auto mixed  precision）——自动混合精度技术在它们的N卡系列上训练模型也构建的一个库。在原有模型训练代码中添加少量代码(4行代码)就能够使用自动混合精度的技术来提高模型的训练速度，提高生产力。简单的理解apex就是一个用来支持模型训练在pytorch框架下使用混合精度进行加速训练的拓展插件之类的库；也可以理解为一种模型训练加速技术——其实是amp自动混合精度搭配硬件一起才能加速的。

amp——auto mixed  precision——自动半精度，它的最核心的东西在于低精度Fp16。它能够提供一种非常可靠和友好的方式进行模型在Fp16精度下进行训练。

2、为什么要使用低精度

深度学习系统大都采用的都是Fp32来进行表示的，随着模型越来越大，以及硬件的进步，加速训练模型的需求就产生了。

现在深度学习系统中广泛使用Fp32进行表示，主要存在2个问题：第一模型尺寸大，训练的时候对显卡的显存要求高；第二模型训练速度慢。

与单精度float（32bit，4个字节）相比，半进度float16仅有16bit，2个字节组成。很明显可以看到，使用Fp16可以解决或者缓解上面fp32的两个问题，fp16的优势就是：

显存占用更少：通用的模型 fp16 占用的内存只需原来的一半，训练的时候可以使用更大的batchsize。

计算速度更快：有论文指出半精度的计算吞吐量可以是单精度的 2-8 倍。

由于fp16 的有效的动态范围比单精度的float要狭窄很多，精度下降(数点后16相比较小数点后8位要精确的多)，那么必然就会产生一系列的问题。如何安全有效的达成使用低精度训练模型显存占用减小和时间减少而不出问题，这里就要使用到一套训练技术和适配的硬件支持了——该技术就是apex提供的amp和N系显卡某些型号。

3、Fp16带来的问题和解决办法

上面说到Fp16会导致一些问题，使得模型训练可能出错，那么到底会有哪些问题呢？

a、溢出错误

Fp16表示的范围是(),Fp32表示的范围是，前者的范围比后者的小很多。当由Fp32计算转化为Fp16的时候，有很大的概率会出现溢出错误。当一个数比还小的时候，Fp16就表示不了了会出现下溢出Underflow；当一个数比65504还要大的时候，就会出现上溢出Overflow。

一般而言，深度学习模型训练过程中，神经网络的激活函数的梯度比权重的梯度要小，也更容易出现下溢出错误。当出现上下出错误的时候，反向传播中误差累积可以把这些数字变成0或者 nans; 这会导致不准确的梯度更新，影响你的网络收敛。

b、舍入误差

 舍入误差指的是当梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败。

知乎——https://zhuanlan.zhihu.com/p/79887894——上找了一张图：

发生舍入误差的时候，权重得不到更新。

针对上述问题，解决的办法是什么呢？聪明的研究人员发明了——混合精度训练+动态损失放大——这种高明的技术来解决上述问题。原理很复杂，简单的描述一下大致思想和步骤：

a、混合精度训练

在内存中用FP16做储存和乘法从而加速计算，用FP32做累加避免舍入误差。这样在权重更新的时候就不会出现舍入误差导致更新失败。

b、损失放大

使用了混合精度训练，还是会存在无法收敛的情况。当梯度值非常小，小到F16不能表示就会出现下溢出错误。就要对损失进行放大：

首先，反向传播之前，把损失值手动增加倍，这个时候反向传播过程中得到的中间变量激活函数梯度之类的不会溢出。

然后，反向传播后，又将权重梯度值缩小倍，这样就解决了下溢出的错误。

当然这里仍然存在着上溢出的问题，那这个时候就直接跳过这一个步的权重更新。只要不是训练过程中一直出现上溢出的问题，训练就是正常的。

二、apex的两种方式

上文中大致描述了半精度训练的一些理论知识，实现层面到底怎么操作呢？首先要有一定的硬件支持——Tensor Core；其次就是要运用相关的代码来实现半精度自动训练。

1、NVIDIA apex

由英伟达开发了一个支持半精度自动训练的pytorch拓展插件，添加几行代码就能实现自动半精度训练。代码流程如下：

from apex import amp
 
>>>>>>>>>>
other code
>>>>>>>>>>
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()#梯度自动缩放
optimizer.step()#优化器更新梯度
>>>>>>>>>>
other code
>>>>>>>>>>

 

这个插件库很好用，解决了一些模型训练显卡显存不足和训练时间长的问题。这里也有一些坑，主要是apex 库安装比较麻烦，不太能顺利安装成功。apex github给出的安装方法：

下载github上的代码和文件，然后通过 pip 安装，不能直接 pip install apex

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

上述安装步骤中，pip install非常容易报错，非常注意的点就是：

pytorch和cuda版本一定要一致；显卡驱动也要适配；还有其他的依赖库najia等也要安装。

2、torch 原生支持的apex

最简单了，只需要安装有pytorch就能使用，而且代码也简单。限制条件只有一个就是pytorch的版本一定>1.6。主要是利用了这两个API——torch.cuda.amp.GradScalar 和 torch.cuda.amp.autocast。

可参考（pytorch_amp: https://pytorch.org/docs/stable/notes/amp_examples.html）

标准流程代码如下：

from torch.cuda.amp import autocast as autocast, GradScaler
 
>>>>>>>>
other code
>>>>>>>>
 
# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()
 
>>>>>>>>
other code
>>>>>>>>
        # 前向过程(model + loss)开启 autocast
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
 
        # Scales loss，这是因为半精度的数值范围有限，因此需要用它放大
        scaler.scale(loss).backward()
 
        # scaler.step() unscale之前放大后的梯度，但是scale太多可能出现inf或NaN
        # 故其会判断是否出现了inf/NaN
        # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
        # 如果检测到出现了inf或者NaN，就跳过这次梯度更新，同时动态调整scaler的大小
        scaler.step(optimizer)
 
        # 查看是否要更新scaler,这个要注意不能丢
        scaler.update()
 
 
>>>>>>>>
other code
>>>>>>>>

三、模型训练效果对比

apex效果到底如何，下面训练一个模型看看实际情况。

并从以下几个方面进行比较：时间，内存，loss趋势，模型准确率

 

参考博文：https://blog.csdn.net/husthy/article/details/109485088