一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)

前言: 本篇博文主要介绍Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)

win 10系统下安装NVIDIA apex

  1. 首先在github下载源码https://github.com/NVIDIA/apex 到本地文件夹
  2. 打开cmd命令窗口,切换到apex所在的文件夹
  3. 使用命令:python setup.py install 即可完成安装

注意事项: 可能会出现的问题:
setuptools有ModuleNotFoundError→更新setuptools

pip install --upgrade setuptools

linux系统下安装NVIDIA apex

安装操作如下所示:

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

Apex的使用教程(PyTorch实现)

首先,Apex是一款由Nvidia开发的基于PyTorch的混合精度训练加速神器,最近Apex更新了API,可以用短短三行代码就能实现不同程度的混合精度加速,训练时间直接缩小一半。可以让你瞬间获得双倍训练速度的快感!想想是是不是很激动。话不多说,直接上代码展示:

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有更深入的了解,下面是一些理论部分的详细讲解:

FP16:半精度浮点数

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

一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)_第1张图片
一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)_第2张图片
一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)_第3张图片

FP16的优点:

  1. 减少显存占用 现在模型越来越大当你使用Bert这一类的预训练模型时,往往显存就被模型及模型计算占去大半,当想要使用更大的Batch Size的时候会显得捉襟见肘。由于FP16的内存占用只有FP32的一半,自然地就可以帮助训练过程节省一半的显存空间。
  2. 加快训练和推断的计算 与普通的空间时间Trade-off的加速方法不同,FP16除了能节约内存,还能同时节省模型的训练时间。在大部分的测试中,基于FP16的加速方法能够给模型训练带来多一倍的加速体验(爽感类似于两倍速看肥皂剧)。
  3. 张量核心的普及 硬件的发展同样也推动着模型计算的加速,随着Nvidia张量核心(Tensor Core)的普及,16bit计算也一步步走向成熟,低精度计算也是未来深度学习的一个重要趋势,再不学习就out啦。

FP16的缺点:量化误差
这个部分是整个博客最重要的理论核心。 讲了这么多FP16的好处,那么使用FP16的时候有没有什么问题呢?当然有。FP16带来的问题主要有两个:1. 溢出错误;2. 舍入误差。
一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)_第4张图片
一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)_第5张图片

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

  1. 混合精度训练(Mixed Precision) 混合精度训练的精髓在于“在内存中用FP16做储存和乘法从而加速计算,用FP32做累加避免舍入误差”。混合精度训练的策略有效地缓解了舍入误差的问题。
  2. 损失放大(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()

一文详解Apex的安装和使用教程(一款基于 PyTorch 的混合精度训练加速神器)_第6张图片

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

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

参考链接

  1. https://zhuanlan.zhihu.com/p/79887894
  2. https://nvidia.github.io/apex/amp.html
  3. https://github.com/NVIDIA/apex
  4. Nvidia官方的混合精度训练文档

你可能感兴趣的:(计算机视觉)