混合精度训练之APEX（上）

跬步亦是千里~~~

——子棐

前言

为了更快的在Pytorch中加入对Volta GPU的支持，并实现针对混合精度训练的优化，NVIDIA发布了Apex开源工具库。

对于混合精度训练问题，该工具库提供了AMP（自动混合精度）以及FP16_Optimizer两种不同的库。AMP提供的是自动的混合精度训练支持，该库会自动检查执行的操作，对于FP16安全的操作在训练中Cast到FP16，反之则选择FP32，从而实现一键混合精度训练，实现了在训练稳定性和易用性的平衡。而FP16_Optimizer提供的是高阶版混合精度训练，其提供了更多详细的实现细节，对于整个网络完全采用FP16，适合最为前沿的玩家更加灵活的实现FP16混合精度训练。

除了针对混合精度训练问题的解决方案，Apex还提供了另外五个库，分别是Apex_C、RNN、Parallel、Reparameterization和Example，这五个库分别是部分功能函数Cuda实现、提供RNN在混合精度训练的实现、提供高性能的多GPU训练（MPI+NCCL）实现、解决模型权值正则化实现问题以及Apex使用范例。

不积跬步无以至千里，因此本文将先从AMP开始介绍，高端玩家可以直接前往FP16_Optimizer区，也欢迎投稿分享经验~~~

Apex之AMP

AMP提供了最为简单的混合精度训练支持，其是通过黑白名单的方式予以实现的。在AMP目录下有Lists目录，其中包括了以下三个库的黑白名单：

Functional库(torch.nn.functional)：其支持FP16的Function为'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d', 'conv_tbc', 'linear'；

tensor库(torch.Tensor和torch.autograd.Variable)：其支持FP16 Fuction为 '__matmul'；

torch库：其支持FP16 Function为：'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d', 'conv_tbc', 'addmm', 'addmv', 'addr', 'matmul', 'mm', 'mv'。

Apex对RNN并不像黑白名单里面直接进行“on-ramp”的cast操作，而是通过rnn_cast函数来支持的，其是通过Wrap torch.nn 后端的顶层的RNN，然后调用CudnnRNN或者AutogradRNN来实现的。

AMP的初始化实现在amp.py中，初始化函数为amp.init()，其内部初始化了一个AmpHandle对象，并对不同的层、函数、RNN进行Wrap配置。在反向迭代时，通过调用AmpHandle函数内部的scale_loss函数进行反向迭代。其中amp.init()函数有三个输入参数，分别是enable、enable_caching和verbose。其中enable是模式打开，即启用amp，如果关闭则在初始化是初始一个NoOpHandle对象，其实不会进行任何的操作。而enable_caching决定是否缓存每次Cast的FP16对象，以解决RNN的问题。而verbose配置打开后可以看到每次的Cast操作。这里以Mnist为例，每次迭代，AMP对整个网络的Cast操作为：

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (linear)

Float->Half (linear)

Float->Half (linear)

Float->Half (linear)

Half->Float (log_softmax)

Apex之Automatic Loss Scaling

关于Loss Scaling已经在前面几期的混合精度训练中有过详细的解释，如果对这个概念还感到陌生的话，可以再次重温来自子棐同学的大作——Tesla V100之NVCaffe混合精度训练深度解析篇（中）。在Apex中或者说Pytorch中所提的Loss Scalling概念和原理和NVcaffe所述一致的，但是在实现上有着不同的思路（如果完全一样这一章节也就没有存在的必要了）。

在Apex中不仅有Loss Scalling而且还有强化升级版“Automatic”，实现算法的如下框所示：

Start with a large scaling factor s

For each training iteration

Make an fp16 copy of weights

Fwd prop

Scale the loss by s

Bwd prop

Update scaling factor s

If dW contains Inf/NaN then reduce s skip the update

If no Inf/NaN were detected for N updates then increase s

Scale dW by 1/s

Update W

在训练开始前会设置一个较大的Scaling Factor S。对于每一次的迭代，和Nvcaffe一样需要维护一个Master权值，因此首先是将FP32权值复制一份用以留存，并生成一份FP16权值。之后利用这个FP16权值进行正向迭代，这个过程就可以充分利用到Tensor Core 125T的计算性能。正向迭代完成后计算Loss，此时的Loss计算值为FP32格式。之后将Loss乘以S以在横轴上右移（放大）Loss获得Scaled_Loss，这样可以在反向迭代转换至FP16时能尽可能的保留FP32的分布特性。下一步基于Scaled_Loss进行反向迭代以计算梯度。计算完成后如果dW中存在着Inf/NaN的话证明存在Overflow，会缩小S值并放弃这一次的更新。如果不存在Inf/NaN则证明放大没有问题，则在N次成功的更新后放大S值，在更新Master权值前要将dW缩小S倍，然后再更新到Master Weight上。

上图是训练一个NMT网络时对S的分析结果图，如图可知，在前面的迭代中S一直在扩大，随后达到最大值，之后一直在一定的范围内上下浮动。通过分析图可知：

最小的S为2的20次方，意味着最大dW不超过2的-5次方；

这种方法能非常高效的为网络设置合适的Scaling Factor。

当然如果是观察仔细的小伙伴肯定会发现几个小问题：

多了一个参数N，作为Scaling Factor的扩大周期；

如果跳过参数更新，是否需要增加迭代计数；

如果跳过参数更新，是否需要对该Batch重新计算。

通过大量的实验总结得知，如果跳过参数更新，是否增加迭代计算并不会对性能造成影响，而且无需重新计算该Batch，直接进行下一个Batch的计算和更新即可。对于参数N，在实验中选取的是2000，而且对于S的缩放采用了2次方的形式。至于上述更新依然可行的原因在于，通过Automatic Loss Scaling可以最高效的实现梯度更新，几个Batch的更新缺失并不会对训练效果造成影响。

本章的最后，给各位小伙伴留一个家庭作业：在Nvcaffe是否有Automatic Loss Scaling呢？如果有，又是怎么实现的呢？

Apex之快速使用指南

安装

软件使用第一步永远都是安装，Apex所需依赖环境为：

Python 3

CUDA 9

PyTorch 0.4

使用Docker Hub提供的最新Cuda9.2的镜像以及PIP安装最新的Pytorch即可正常编译安装Apex。

更加推荐使用NGC已经优化的容器镜像，使用方法如下：

前往ngc.nvidia.com注册登录并获取API Key

使用API Key登录到容器镜像仓库，命令：docker login

拉取最新容器镜像：docker pull nvcr.io/nvidia/pytorch:18.05-py3

运行容器：docker run --runtime=nvidia --rm -tinvcr.io/nvidia/pytorch:18.05-py3 bash

下载Apex源码：git clone https://github.com/NVIDIA/apex.git

编译安装Apex：cd apex && python setup.py install

*这里需要注意的是NGC的Pytorch建议使用18.04之后版本。

针对单一Optimizer使能AMP

对于单一Optimizer，使能AMP只需要以下修改两部分即可：

1. 代码中Optimizer之后添加：

optimizer = # ... some optimizer

from apex import amp

amp_handle = amp.init()

2. 修改代码中反向迭代部分：

# ... do a bunch of stuff to compute a loss

loss.backward()

optimizer.step()

# ...finish the iteration

修改为

# ... same as before

with amp_handle.scale_loss(loss, optimizer) as scaled_loss:

scaled_loss.backward()

optimizer.step()

# ... same as before

针对多Optimizer使能AMP

多Optimizer使能AMP相较单一Optimizer复杂，需要显性的warp多个Optimizer，并处理多个Loss的反向迭代，例如：

amp_handle = amp.init()

optimizer = amp_handle.wrap_optimizer(optimizer, num_loss=2)

# ...

optimizer.zero_grad()

loss1 = ComputeLoss1(model)

with optimizer.scale_loss(loss1) as scaled_loss:

scaled_loss.backward()

# ...

loss2 = ComputeLoss2(model)

with optimizer.scale_loss(loss2) as scaled_loss:

scaled_loss.backward()

# ...

optimizer.step()

AMP添加自定义函数

在AMP添加自定义函数的支持有两种方法，分别使用Function annotation或者函数显性注册如下所示：

对于自定义函数Fru

from backend import FRUBackend

def fru(input, hidden, weight, bias):

# call to CUDA code

FRUBackend(input, hidden, weight, bias)

使用Function Annotation

@amp.half_function

def fru(input, hidden, weight, bias):

#...

使用显性注册

import backend

amp.register_half_function(backend, 'FRUBackend')

amp.init()

此处的Function Annotation有三种类型分别是：

类似显性注册也有对应三种方式，分别是：

amp.register_half_function(module, function_name)

amp.register_float_function(module, function_name)

amp.register_promote_function(module, function_name)

这里需要注意的是，采用显性注册需要在amp.init()之前。

至此，关于AMP的介绍和使用指南已经完全阐明，关于前文留的小问题，还欢迎小伙伴们留言回答哟~~~

以上

子棐