【PyTorch】加速不香吗？| 基于Apex的混合精度加速

会员服务 ·

【PyTorch】加速不香吗？| 基于Apex的混合精度加速

2020 年 5 月 1 日 计算机视觉战队

点击蓝字关注我们

扫码关注我们

公众号 : 计算机视觉战队

加入我们，大量论文代码下载链接

背景

我们提到圆周率 π 的时候，它有很多种表达方式，既可以用数学常数3.14159表示，也可以用一长串1和0的二进制长串表示。

圆周率 π 是个无理数，既小数位无限且不循环。因此，在使用圆周率进行计算时，人和计算机都必须根据精度需要将小数点后的数字四舍五入。

在小学的时候，小学生们可能只会用手算的方式计算数学题目，圆周率的数值也只能计算到小数点后两位——3.14；而高中生使用图形计算器可能会使圆周率数值排到小数点后10位，更加精确地表示圆周率。在计算机科学中，这被称为精度，它通常以二进制数字来衡量，而非小数。

对于复杂的科学模拟，开发人员长期以来一直都依靠高精度数学来研究诸如宇宙大爆炸，或是预测数百万个原子之间的相互作用。

数字位数越高，或是小数点后位数越多，意味着科学家可以在更大范围内的数值内体现两个数值的变化。借此，科学家可以对最大的星系，或是最小的粒子进行精确计算。

但是，计算精度越高，意味着所需的计算资源、数据传输和内存存储就越多。其成本也会更大，同时也会消耗更多的功率。

由于并非每个工作负载都需要高精度，因此 AI 和 HPC 研究人员可以通过混合或匹配不同级别的精度的方式进行运算，从而使效益最大化。NVIDIA Tensor Core GPU 支持多精度和混合精度技术，能够让开发者优化计算资源并加快 AI 应用程序及其推理功能的训练。

单精度、双精度和半精度浮点格式之间的区别

IEEE 浮点算术标准是用来衡量计算机上以二进制所表示数字精度的通用约定。在双精度格式中，每个数字占用64位，单精度格式占用32位，而半精度仅16位。

要了解其中工作原理，我们可以拿圆周率举例。在传统科学记数法中，圆周率表示为3.14 x100。但是计算机将这些信息以二进制形式存储为浮点，即一系列的1和0，它们代表一个数字及其对应的指数，在这种情况下圆周率则表示为1.1001001 x 21。

在单精度32位格式中，1位用于指示数字为正数还是负数。指数保留了8位，这是因为它为二进制，将2进到高位。其余23位用于表示组成该数字的数字，称为有效数字。

而在双精度下，指数保留11位，有效位数为52位，从而极大地扩展了它可以表示的数字范围和大小。半精度则是表示范围更小，其指数只有5位，有效位数只有10位。

圆周率在每个精度级别表现如下：

多精度和混合精度计算的差异

多精度计算意味着使用能够以不同精度进行计算的处理器，在需要使用高精度进行计算的部分使用双精度，并在应用程序的其他部分使用半精度或单精度算法。

混合精度（也称为超精度）计算则是在单个操作中使用不同的精度级别，从而在不牺牲精度的情况下实现计算效率。

在混合精度中，计算从半精度值开始，以进行快速矩阵数学运算。但是随着数字的计算，机器会以更高的精度存储结果。例如，如果将两个16位矩阵相乘，则结果为32位大小。

使用这种方法，在应用程序结束计算时，其累积得到结果，在准确度上可与使用双精度算法运算得到的结果相媲美。

这项技术可以将传统的双精度应用程序加速多达25倍，同时减少了运行所需的内存、时间和功耗。它可用于 AI 和模拟 HPC 工作负载。

随着混合精度算法在现代超级计算应用程序中的普及，HPC 专家 Jack Dongarra 提出了一个新的基准，即 HPL-AI，以评估超级计算机在混合精度计算上的性能。当 NVIDIA 在全球最快的超级计算机 Summit 上运行 HPL-AI 计算时，该系统达到了前所未有的性能水平，接近550 petaflop，比其在 TOP500 榜单上的官方性能快了3倍。

理论

PyTorch实现

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

但是如果你希望对FP16和Apex有更深入的了解，或是在使用中遇到了各种不明所以的“Nan”的同学，可以接着读下去，后面会有一些有趣的理论知识和瓦砾最近一个月使用Apex遇到的各种bug，不过当你深入理解并解决掉这些bug后，你就可以彻底摆脱“慢吞吞”的FP32啦。

理论部分

为了充分理解混合精度的原理，以及API的使用，先补充一点基础的理论知识。

什么是FP16？

半精度浮点数是一种计算机使用的二进制浮点数数据类型，使用2字节（16位）存储。

其中，sign位表示正负，exponent位表示指数（$2^{n-15+{1}(n=0)}$），fraction位表示的是分数（$\frac{m}{1024}$）。其中当指数为零的时候，下图加号左边为0，其他情况为1。

FP16的表示范例

为什么需要FP16？

在使用FP16之前，我想再赘述一下为什么我们使用FP16。

减少显存占用

现在模型越来越大，当你使用Bert这一类的预训练模型时，往往显存就被模型及模型计算占去大半，当想要使用更大的Batch Size的时候会显得捉襟见肘。由于FP16的内存占用只有FP32的一半，自然地就可以帮助训练过程节省一半的显存空间。

加快训练和推断的计算

与普通的空间时间Trade-off的加速方法不同，FP16除了能节约内存，还能同时节省模型的训练时间。在大部分的测试中，基于FP16的加速方法能够给模型训练带来多一倍的加速体验（爽感类似于两倍速看肥皂剧）。

张量核心的普及

硬件的发展同样也推动着模型计算的加速，随着Nvidia张量核心（Tensor Core）的普及，16bit计算也一步步走向成熟，低精度计算也是未来深度学习的一个重要趋势，再不学习就out啦。

FP16带来的问题：量化误差

这个部分是整个博客最重要的理论核心。讲了这么多FP16的好处，那么使用FP16的时候有没有什么问题呢？当然有。FP16带来的问题主要有两个：1. 溢出错误；2. 舍入误差。

溢出错误（Grad Overflow / Underflow）

由于FP16的动态范围（$6 \times 10^{-8} \sim 65504$）比FP32的动态范围（$1.4 \times 10^{-45} \sim 1.7 \times 10^{38}$）要狭窄很多，因此在计算过程中很容易出现上溢出（Overflow，$g>65504$）和下溢出（Underflow，$g<6\times10^{-8}$）的错误，溢出之后就会出现“Nan”的问题。

在深度学习中，由于激活函数的的梯度往往要比权重梯度小，更易出现下溢出的情况。

舍入误差（Rounding Error）

舍入误差指的是当梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败，用一张图清晰地表示：

解决问题的办法：混合精度训练+动态损失放大

混合精度训练（Mixed Precision）

混合精度训练的精髓在于“在内存中用FP16做储存和乘法从而加速计算，用FP32做累加避免舍入误差”。混合精度训练的策略有效地缓解了舍入误差的问题。

损失放大（Loss Scaling）

即使用了混合精度训练，还是会存在无法收敛的情况，原因是激活梯度的值太小，造成了下溢出（Underflow）。损失放大的思路是：

反向传播前，将损失变化（dLoss）手动增大$2^k$倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；
反向传播后，将权重梯度缩$2^k$倍，恢复正常值。

Apex的新API：Automatic Mixed Precision (AMP)

曾经的Apex混合精度训练的api仍然需要手动half模型已经输入的数据，比较麻烦，现在新的api只需要三行代码即可无痛使用：

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

opt_level

其中只有一个 opt_level 需要用户自行配置：

O0 ：纯FP32训练，可以作为accuracy的baseline；
O1 ：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算。
O2 ：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。
O3 ：纯FP16训练，很不稳定，但是可以作为speed的baseline；

动态损失放大（Dynamic Loss Scaling）

AMP默认使用动态损失放大，为了充分利用FP16的范围，缓解舍入误差，尽量使用最高的放大倍数（$2^{24}$），如果产生了上溢出（Overflow），则跳过参数更新，缩小放大倍数使其不溢出，在一定步数后（比如2000步）会再尝试使用大的scale来充分利用FP16的范围：

踩过的坑

这一部分是整篇最干货的部分，在最近在apex使用中的踩过的所有的坑，由于apex报错并不明显，常常debug得让人很沮丧，但只要注意到以下的点，95%的情况都可以畅通无阻了：

判断你的GPU是否支持FP16：构拥有Tensor Core的GPU（2080Ti、Titan、Tesla等），不支持的（Pascal系列）就不建议折腾了。
常数的范围：为了保证计算不溢出，首先要保证人为设定的常数（包括调用的源码中的）不溢出，如各种epsilon，INF等。
Dimension最好是8的倍数： Nvidia官方的文档的2.2条表示，维度都是8的倍数的时候，性能最好。
涉及到sum的操作要小心，很容易溢出，类似Softmax的操作建议用官方API，并定义成layer写在模型初始化里。
模型书写要规范：自定义的Layer写在模型初始化函数里，graph计算写在forward里。
某些不常用的函数，在使用前需要注册：amp.register_float_function(torch, 'sigmoid')
某些函数（如einsum）暂不支持FP16加速，建议不要用的太heavy， xlnet的实现改FP16 困扰了我很久。
需要操作模型参数的模块（类似EMA），要使用AMP封装后的model。
需要操作梯度的模块必须在optimizer的step里，不然AMP不能判断grad是否为Nan。

总结

这篇从理论到实践地介绍了混合精度计算以及Apex新API（AMP）的使用方法。现在在做深度学习模型的时候，几乎都会第一时间把代码改成混合精度训练的了，速度快，精度还不减，确实是调参炼丹必备神器。目前网上还并没有看到关于AMP以及使用时会遇到的坑的中文博客，所以这一篇也是希望大家在使用的时候可以少花一点时间debug。

如果想加入我们“ 计算机视觉研究院 ”，请扫二维码加入我们。我们会按照你的需求将你拉入对应的学习群！

计算机视觉研究院主要涉及深度学习领域，主要致力于人脸检测、人脸识别，多目标检测、目标跟踪、图像分割等研究方向。研究院接下来会不断分享最新的论文算法新框架，我们这次改革不同点就是，我们要着重” 研究 “。之后我们会针对相应领域分享实践过程，让大家真正体会 摆脱理论 的真实场景，培养爱动手编程爱动脑思考的习惯！