pytorch原生支持的amp混合精度和nvidia apex混合精度AMP技术

一. apex原理简介

1. apex和amp

从英伟达网页Apex (A PyTorch Extension) — Apex 0.1.0 documentation可以得到apex的全称——A PyTorch Extension(Apex)——其实就是一种pytorch的拓展插件。它的目的就是为了是用户能够快速实现amp——自动混合精度技术在它们的N卡系列上训练模型也构建的一个库。简简单单的在你自定义的模型训练代码中添加少量代码(4行代码)就能够使用自动混合精度的技术来提高模型的训练速度，提高生产力。简单的理解apex就是一个用来支持模型训练在pytorch框架下使用混合精度进行加速训练的拓展插件之类的库；也可以理解为一种模型训练加速技术——其实是amp自动混合精度搭配硬件一起才能加速的。

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

2. 为什么要使用低精度

深度学习系统大都采用的都是Fp32来进行表示的，随着模型越来越大，已经硬件的进步，加速训练模型的需求就产生了。现在深度学习系统中广泛使用Fp32进行表示，主要存在2个问题，第一模型尺寸大，训练的时候对显卡的显存要求高；第二模型训练速度慢。

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

1. 显存占用更少：通用的模型 fp16 占用的内存只需原来的一半，训练的时候可以使用更大的batchsize。
2. 计算速度更快：有论文指出半精度的计算吞吐量可以是单精度的 2-8 倍。

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

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

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

3.1 溢出错误

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

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

3.2 舍入误差

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

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

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

1. 混合精度训练：
   在内存中用FP16做储存和乘法从而加速计算，用FP32做累加避免舍入误差。这样在权重更新的时候就不会出现舍入误差导致更新失败。
2. 损失放大：
   使用了混合精度训练，还是会存在无法收敛的情况。当梯度值非常小，小到F16不能表示就会出现下溢出错误。就要对损失进行放大：
   首先，反向传播之前，把损失变化手动增加倍，这个时候反向传播过程中得到的中间变量激活函数梯度之类的不会溢出。
   然后，反向传播后，又将权重梯度值缩小倍，这样就解决了下溢出的错误。
   当然这里仍然存在着上溢出的问题，那这个时候就直接跳过这一个步的权重更新。只要不是训练过程中一直出现上溢出的问题，训练就是正常的。

二. amp两种方式

上文中大致描述了半精度训练的一些理论知识，实现层面到底怎么操作呢？首先要有一定的硬件支持——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()#优化器更新梯度
optimizer.zero_grad()
>>>>>>>>>>
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 -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

非常容易报错（如下面的报错信息），非常注意的点就是：
pytorch和cuda版本一定要一致；显卡驱动也要适配；还有其他的依赖库najia等也要安装。

注意：如果上面安装步骤中的第三行命令报错：

报错内容：RuntimeError: Cuda extensions are being compiled with a version of Cuda that does not match the version used to compile Pytorch binaries. Pytorch binaries were compiled with Cuda 9.0.176.

解决方式PyTorch中apex安装方式和避免踩坑：其错误意思就是cuda和pytorch的版本不对应，但是通过搜索也发现可以不带 --global --option 也能用
于是，修改第三行命令为：

pip install -v --no-cache-dir ./

2. torch 原生支持的amp

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

注意：Gradscalar 需要对梯度更新计算(检查是否溢出)和优化器(将丢弃的batches转换为 no-op)进行控制，以实现其操作。 这就是为什么 loss.backwards()被 scaler.scale(loss).backwards()取代， 以及 optimizer.step()被 scaler.step(optimizer)替换的原因。

标准流程代码如下：

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
>>>>>>>>

三. 具体模型训练时使用例子

1. train_without_amp：模型无加速情况

    for epoch in tqdm(range(args.epochs), desc='Epoch'):
        pbar = ProgressBar(n_total=len(train_dataloader), desc='Training Iteraction')
        for step, batch in enumerate(train_dataloader):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            inputs = {'input_ids': batch[0], 'attention_mask': batch[1], 'bio_labels': batch[2]}
            model.zero_grad()
            optimizer.zero_grad()
            outputs = model(**inputs)
            loss = outputs[0]
            loss.backward()
            optimizer.step()
            scheduler.step()
            total_loss += loss.item()
            losses.append(loss.item())
            pbar(step, {"loss": loss.item()})
            memo_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            gpu_use_rato.append(round(memo_info.used/1024/1024/1024,4))

2. train_torch_amp：使用pytorch自带的加速amp

    scaler = GradScaler()
    gpu_use_rato = []
    losses = []
    t1 = time.time()
    for epoch in tqdm(range(args.epochs), desc='Epoch'):
        pbar = ProgressBar(n_total=len(train_dataloader), desc='Training Iteraction')
        for step, batch in enumerate(train_dataloader):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            inputs = {'input_ids': batch[0], 'attention_mask': batch[1], 'bio_labels': batch[2]}
            model.zero_grad()
            optimizer.zero_grad()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            with autocast():
                outputs = model(**inputs)
                loss = outputs[0]
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            scaler.step(optimizer)
            scaler.update()
            total_loss += loss.item()
            losses.append(loss.item())
            pbar(step, {"loss": loss.item()})
            memo_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            gpu_use_rato.append(round(memo_info.used/1024/1024/1024,4))

3. train_nvidia_amp：使用nvidia 的加速amp

    model,optimizer = amp.initialize(model,optimizer,opt_level='O1')
    for epoch in tqdm(range(args.epochs), desc='Epoch'):
        pbar = ProgressBar(n_total=len(train_dataloader), desc='Training Iteraction')
        for step, batch in enumerate(train_dataloader):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            inputs = {'input_ids': batch[0], 'attention_mask': batch[1], 'bio_labels': batch[2]}
            outputs = model(**inputs)
            loss = outputs[0]
            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
            model.zero_grad()
            optimizer.zero_grad()
 
            with amp.scale_loss(loss,optimizer) as scaled_loss:
                scaled_loss.backward()
            optimizer.step()
            total_loss += loss.item()
            losses.append(loss.item())
            scheduler.step()
            pbar(step, {"loss": loss.item()})
            memo_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            gpu_use_rato.append(round(memo_info.used/1024/1024/1024,4)) 

四. 相关链接

pytorch原生支持的apex混合精度和nvidia apex混合精度AMP技术加速模型训练效果对比
训练提速60%！只需5行代码，PyTorch 1.6即将原生支持自动混合精度训练。