一行代码让训练速度提升2倍，飞桨自动混合精度技术详解

机器之心编辑部

飞桨自动混合精度技术，让你的训练速度飞起来。

随着生活节奏的加快，「等待」已经越来越成为人们希望远离的事情。但是在深度学习领域，模型的参数、数据集的规模等等动辄就是以亿为单位，甚至更大，因此当模型训练成功之时，放一首张靓颖的「终于等到你」作为背景音乐实在是太应景了。

那如果现在向你推荐一款神器，可以实现训练速度翻倍，访存效率翻倍，你心动吗？心动不如行动（这可不是电视直销，别着急换频道），来和我一起看看这款神器——基于飞桨核心框架的自动混合精度（Automatic Mixed Precision) 技术，简称飞桨 AMP 技术。

飞桨 AMP 技术仅仅通过一行代码即可帮助用户简便快速的将单精度训练的模型修改为自动混合精度训练。同时通过黑白名单和动态 Loss Scaling 来保证训练的稳定性，避免出现 INF 或者 NAN 问题。飞桨 AMP 可以充分发挥新一代 NVIDIA GPU 中 Tensor Core 的计算性能优势，ResNet50、Transformer 等模型的训练速度与单精度训练相比可以提升到 1.5～2.9 倍。

那么它是怎么实现的呢？我们先从什么是自动混合精度技术讲起。

什么是自动混合精度技术

顾名思义，自动混合精度是一种自动将半精度和单精度混合使用，从而加速模型训练的技术。其中单精度（Float Precision32，FP32）好理解，是计算机常用的一种数据类型。那么半精度是什么呢？如图 1 所示，半精度（Float Precision16，FP16）是一种相对较新的浮点类型，在计算机中使用 2 字节（16 位）存储，在 IEEE 754-2008 中，它被称作 binary16。与计算中常用的单精度和双精度类型相比，Float16 更适于在精度要求不高的场景中使用。

图 1 半精度和单精度数据示意图

不言而喻，在深度学习领域，如果使用 Float16 代替 Float32 来存储数据，那么开发者就可以训练更大更复杂的模型，使用更大的 batch size。因此对于那些恨不得挖掘出 GPU 里每一个晶体管全部潜力的科学家们怎么能放过它呢？同时由于 NVIDIA 推出了具备 Tensor Core 技术的 Volta 及 Turing 架构 GPU，使半精度计算趋向成熟。在相同的 GPU 硬件上，Tensor Core 的半精度计算吞吐量是单精度的 8 倍。

但显而易见，使用 Float16 肯定会同时带来计算精度上的损失。但对深度学习训练而言，并不是所有计算都要求很高的精度，一些局部的精度损失对最终训练效果影响很微弱，仅需要某些特殊步骤保留 Float32 的计算精度即可。因此混合精度计算的需求应运而生。我们可以将训练过程中一些对精度损失不敏感且能使用 Tensor Core 进行加速的运算使用半精度处理，最大限度的提升访存和计算效率。

但是对每个具体模型，人工去设计和尝试精度混合的方法，是非常繁琐的，我们迫切需要一种更简洁的方式，高效地实现混合精度的训练。AMP，顾名思义，就是让混合精度训练自动化，因此使用简单是它的重要特色。具体咋用，咱们往下看！

AMP 的使用方法

下面以 MNIST 为例介绍如何使用飞桨 AMP 技术。MNIST 网络定义的代码如下所示。其中 conv2d、batch_norm（bn）和 pool2d 的数据布局需要提前设置为'NHWC'，这样有利于加速混合精度训练，并且 conv2d 的输出通道数需要设置为 4 的倍数，以便使用 Tensor Core 技术加速。

import paddle.fluid as fluiddef MNIST(data, class_dim):     conv1 = fluid.layers.conv2d(data, 16, 5, 1, act=None, data_format='NHWC')     bn1 = fluid.layers.batch_norm(conv1, act='relu', data_layout='NHWC')     pool1 = fluid.layers.pool2d(bn1, 2, 'max', 2, data_format='NHWC')     conv2 = fluid.layers.conv2d(pool1, 64, 5, 1, act=None, data_format='NHWC')     bn2 = fluid.layers.batch_norm(conv2, act='relu', data_layout='NHWC')     pool2 = fluid.layers.pool2d(bn2, 2, 'max', 2, data_format='NHWC')     fc1 = fluid.layers.fc(pool2, size=50, act='relu')     fc2 = fluid.layers.fc(fc1, size=class_dim, act='softmax')    return fc2

为了训练 MNIST 网络，还需要定义损失函数来更新权重参数，此处使用的优化损失函数是 SGDOptimizer。为了简化说明，这里省略了迭代训练的相关代码，仅体现损失函数及优化器定义相关的内容。

import paddle.fluid as fluidimport numpy as npdata = fluid.layers.data(     name='image', shape=[None, 28, 28, 1], dtype='float32')label = fluid.layers.data(name='label', shape=[None, 1], dtype='int64')out = MNIST(data, class_dim=10)loss = fluid.layers.cross_entropy(input=out, label=label)avg_loss = fluid.layers.mean(loss)sgd = fluid.optimizer.SGDOptimizer(learning_rate=1e-3)sgd.minimize(avg_loss)

那么如何将上面的示例改造成使用 AMP 训练的方式呢？用户仅需要使用飞桨提供的 AMP 函数 fluid.contrib.mixed_precision.decorate 将原来的优化器 SGDOptimizer 进行封装，然后使用封装后的优化器（mp_sgd）更新参数梯度，代码如下所示：

sgd = fluid.optimizer.SGDOptimizer(learning_rate=1e-3)mp_sgd = fluid.contrib.mixed_precision.decorator.decorate(sgd)mp_sgd.minimize(avg_loss)

如上即为最简单的飞桨 AMP 功能使用方法。

但是大家可能有些疑问，模型是如何感知哪些算子（Op）需要被转换呢？是不是还需要手工指定呢？算子那么多，我怎么知道哪个算子可以被转换呢？别着急，飞桨已经帮你定制好了，这也是这门技术被称为「自动」的原因之一，且请往下看！

黑白名单功能

为了让开发者可以方便快捷的使用混合精度计算，飞桨的工程师们使用了大量模型在不同应用场景中反复验证，然后根据半精度数据类型计算的稳定性和加速效果，梳理出一系列适合转换为半精度计算的算子，并将这些算子定义到了一份白名单文件中。同时对于一些经过验证发现不适合转换的算子，也就是使用半精度计算会导致数值不精确的算子将被记录到黑名单文件中。此外一些对半精度计算没有多少影响的算子归类于灰名单。在使用 AMP 训练过程中，系统会自动读取黑白名单，从而感知到哪些算子需要被转换为半精度计算。

对于某些特殊场景，如果开发者希望使用自定义的黑白名单，则可以使用 AutoMixedPrecisionLists 类设置，代码示例如下所示。

sgd = SGDOptimizer(learning_rate=1e-3)# 指定自定义的黑白名单，其中 list1 和 list2 为包含有算子名称的列表amp_list = AutoMixedPrecisionLists(custom_white_list=list1,custom_black_list=list2)mp_sgd = fluid.contrib.mixed_precision.decorator.decorate(sgd, amp_list)mp_sgd.minimize(avg_loss)

那么自动混合精度技术被称为「自动」的原因之二呢？那就是下面的自动调整 Loss Scaling 功能。

自动调整 Loss Scaling

AMP 技术在提升访存和计算效率的同时，伴随的副作用也是很明显的。那就是由于半精度数据类型的精度范围与转换前的单精度相比过窄，导致容易产生 INF 和 NAN 问题。为了避免此类问题，AMP 技术实现了自动调整 Loss Scaling 功能，即在 AMP 训练过程中，为了避免精度下溢，每训练一定数量批次的数据，就将 Loss 放大指定倍数。如果 Loss 在放大过程中发生上溢，则可以再缩小一定倍数，确保整个训练过程中，梯度可以正常收敛。

fluid.contrib.mixed_precision.decorate 函数携带了自动调整 Loss Scaling 功能相关的参数，这些参数都带有默认值，如下面代码所示。这些默认值都是经过飞桨工程师多次验证后定义的。通常情况下，用户可以直接使用，无需重新设置。

sgd = SGDOptimizer(learning_rate=1e-3)mp_sgd = fluid.contrib.mixed_precision.decorator.decorate(sgd,             init_loss_scaling=2**15,             incr_every_n_steps=2000,             use_dynamic_loss_scaling=True)mp_sgd.minimize(avg_loss)

多卡 GPU 训练的优化

在新发布的飞桨核心框架 1.7 版本上，AMP 技术深度优化了多卡 GPU 训练。如图 2 所示，在优化之前的参数梯度更新过程中，梯度计算时虽然使用的是半精度数据类型，但是不同 GPU 卡之间的梯度传输数据类型仍为单精度。

图 2 1.7 版本之前的参数梯度更新过程示意图

为了降低 GPU 多卡之间的梯度传输带宽，我们将梯度传输这个过程提到 Cast 操作之前，而每个 GPU 卡在得到对应的半精度梯度后再执行 Cast 操作，将其转变为单精度类型，如图 3 所示。这一优化在训练网络复杂度较大的模型时，对减少带宽占用方面非常有效，如多卡训练 BERT-Large 模型。

图 3 1.7 版本的参数梯度更新过程示意图

训练性能对比（AMP VS FP32）

飞桨 AMP 技术在 ResNet50、Transformer 等模型上训练速度相对于 FP32 训练来说有非常大的优势，下面以 ResNet50 模型为例，从下图中可以看出，ResNet50 的 AMP 训练相对与 FP32 训练，单卡加速比可达 2.9 倍，八卡加速比可达 2.8 倍。