混合精度

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

概述

混合精度是指训练时在模型中同时使用 16 位和 32 位浮点类型,从而加快运行速度,减少内存使用的一种训练方法。通过让模型的某些部分保持使用 32 位类型以保持数值稳定性,可以缩短模型的单步用时,而在评估指标(如准确率)方面仍可以获得同等的训练效果。本文介绍如何使用实验性 Keras 混合精度 API 来加快模型速度。利用此 API 可以在现代 GPU 上将性能提高三倍以上,而在 TPU 上可以提高 60%。

注:Keras 混合精度 API 目前是实验版本,可能会更改。

如今,大多数模型使用 float32 dtype,这种数据类型占用 32 位内存。但是,还有两种精度较低的 dtype,即 float16 和 bfloat16,它们都是占用 16 位内存。现代加速器使用 16 位 dtype 执行运算的速度更快,因为它们有执行 16 位计算的专用硬件,并且从内存中读取 16 位 dtype 的速度也更快。

NVIDIA GPU 使用 float16 执行运算的速度比使用 float32 快,而 TPU 使用 bfloat16 执行运算的速度也比使用 float32 快。因此,在这些设备上应尽可能使用精度较低的 dtype。但是,出于对数值的要求,为了让模型训练获得相同的质量,一些变量和计算仍需使用 float32。利用 Keras 混合精度 API,float16 或 bfloat16 可以与 float32 混合使用,从而既可以获得 float16/bfloat16 的性能优势,也可以获得 float32 的数值稳定性。

注:在本指南中,术语“数值稳定性”是指使用较低精度的 dtype(而不是较高精度的 dtype)对模型质量的影响。如果使用 float16 或 bfloat16 执行运算,则与使用 float32 执行运算相比,使用这些较低精度的 dtype 会导致模型获得的计算准确率或其他指标相对较低,那么我们就说这种运算“数值不稳定”。

设置

TensorFlow 2.1 中提供了 Keras 混合精度 API。

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.mixed_precision import experimental as mixed_precision

支持的硬件

虽然混合精度在大多数硬件上都可以运行,但是在最新的 NVIDIA GPU 和 Cloud TPU 上才能加速模型。NVIDIA GPU 支持混合使用 float16 和 float32,而 TPU 则支持混合使用 bfloat16 和 float32。

在 NVIDIA GPU 中,计算能力为 7.0 或更高的 GPU 可以获得混合精度的最大性能优势,因为这些型号具有称为 Tensor 核心的特殊硬件单元,可以加速 float16 矩阵乘法和卷积运算。旧款 GPU 使用混合精度无法实现数学运算性能优势,不过可以节省内存和带宽,因此也可以在一定程度上提高速度。您可以在 NVIDIA 的 CUDA GPU 网页上查询 GPU 的计算能力。可以最大程度从混合精度受益的 GPU 示例包括 RTX GPU、Titan V 和 V100。

注:如果在 Google Colab 中运行本指南中示例,则 GPU 运行时通常会连接 P100。P100 的计算能力为 6.0,预计速度提升不明显。

您可以使用以下命令检查 GPU 类型。如果要使用此命令,必须安装 NVIDIA 驱动程序,否则会引发错误。

nvidia-smi -L
GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-d0ab7e62-1881-971d-1d9e-2c37cb2efad2)

所有 Cloud TPU 均支持 bfloat16。

即使在预计无法提升速度的 CPU 和旧款 GPU 上,混合精度 API 仍可以用于单元测试、调试,或试用 API​​。

设置 dtype 策略

要在 Keras 中使用混合精度,您需要创建一条 tf.keras.mixed_precision.experimental.Policy,通常将其称为 dtype 策略。Dtype 策略可以指定将在其中运行的 dtype 层。在本指南中,您将从字符串 'mixed_float16' 构造策略,并将其设置为全局策略。这会导致随后创建的层使用 float16 和 float32 的混合精度。

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)
INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: Tesla V100-SXM2-16GB, compute capability 7.0

该策略指定了层的两个重要方面:完成层的计算所使用的 dtype 和层变量的 dtype。上面的代码创建了一条 mixed_float16 策略(即通过将字符串 'mixed_float16' 传递给其构造函数而构建的 mixed_precision.Policy )。凭借此策略,层可以使用 float16 计算和 float32 变量。计算使用 float16 来提高性能,而变量使用 float32 来确保数值稳定性。您可以直接在策略中查询这些属性。

print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)
Compute dtype: float16
Variable dtype: float32

如前所述,在计算能力至少为 7.0 的 NVIDIA GPU 上,mixed_float16 策略可以大幅提升性能。在其他 GPU 和 CPU 上,该策略也可以运行,但可能无法提升性能。对于 TPU,则应使用 mixed_bfloat16 策略。

构建模型

接下来,我们开始构建一个简单的模型。过小的模型往往无法获得混合精度的优势,因为 TensorFlow 运行时的开销通常占据大部分执行时间,导致 GPU 的性能提升几乎可以忽略不计。因此,如果使用 GPU,我们会构建两个比较大的 Dense 层,每个层具有 4096 个单元。

inputs = keras.Input(shape=(784,), name='digits')
if tf.config.list_physical_devices('GPU'):
  print('The model will run with 4096 units on a GPU')
  num_units = 4096
else:
  # Use fewer units on CPUs so the model finishes in a reasonable amount of time
  print('The model will run with 64 units on a CPU')
  num_units = 64
dense1 = layers.Dense(num_units, activation='relu', name='dense_1')
x = dense1(inputs)
dense2 = layers.Dense(num_units, activation='relu', name='dense_2')
x = dense2(x)
The model will run with 4096 units on a GPU

每个层都有一条策略,默认情况下会使用全局策略。因此,每个 Dense 层都具有 mixed_float16 策略,这是因为之前已将 mixed_float16 设置为全局策略。这样,dense 层就会执行 float16 计算,并使用 float32 变量。为了执行 float16 计算,它们会将输入转换为 float16 类型,因此,输出也是 float16 类型。它们的变量是 float32 类型,在调用层时,它们会将变量转换为 float16 类型,从而避免 dtype 不匹配所引起的错误。

print('x.dtype: %s' % x.dtype.name)
# 'kernel' is dense1's variable
print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)
x.dtype: float16
dense1.kernel.dtype: float32

接下来创建输出预测。通常,您可以按如下方法创建输出预测,但是对于 float16,其结果不一定具有数值稳定性。

# INCORRECT: softmax and model output will be float16, when it should be float32
outputs = layers.Dense(10, activation='softmax', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)
Outputs dtype: float16

模型末尾的 softmax 激活值本应为 float32 类型。但由于 dtype 策略是 mixed_float16,softmax 激活通常会使用 float16 dtype 进行计算,并且会输出 float16 张量。

这一问题可以通过分离 Dense 和 softmax 层,并将 dtype='float32' 传递至 softmax 层来解决。

# CORRECT: softmax and model output are float32
x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)
Outputs dtype: float32

dtype='float32' 传递至 softmax 层的构造函数会将该层的 dtype 策略重写为 float32 策略,从而由后者执行计算并保持变量为 float32 类型。同样,我们也可以传递 dtype=mixed_precision.Policy('float32');层始终将 dtype 参数转换为策略。由于 Activation 层没有变量,因此会忽略该策略的变量 dtype,但是该策略的计算 dtype 为 float32,因此 softmax 和模型的输出也是 float32。

您可以在模型中间添加 float16 类型的 softmax,但模型末尾的 softmax 应为 float32 类型。原因是,如果从 softmax 传递给损失函数的中间张量是 float16 或 bfloat16 类型,则会出现数值问题。

如果您认为使用 float16 计算无法获得数值稳定性,则可以通过传递 dtype='float32',将任何层的 dtype 重写为 float32 类型。但通常,只有模型的最后一层才需要这样重写,因为对大多数层来说,mixed_float16mixed_bfloat16 的精度已经足够。

即使模型不以 softmax 结尾,输出也仍是 float32。虽然对这一特定模型来说并非必需,但可以使用以下代码将模型输出转换为 float32 类型:

# The linear activation is an identity function. So this simply casts 'outputs'
# to float32. In this particular case, 'outputs' is already float32 so this is a
# no-op.
outputs = layers.Activation('linear', dtype='float32')(outputs)

接下来完成并编译该模型,然后生成输入数据。

model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(loss='sparse_categorical_crossentropy',
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255

本示例将输入数据从 int8 强制转换为 float32。我们不转换为 float16 是因为在 CPU 上除以 255 时,float16 的运算速度比 float32 慢。在这种情况下,性能差距可以忽略不计,但一般来说,在 CPU 上执行运算时,数学处理输入应使用 float32 类型。该模型的第一层会将输入转换为 float16,因为每一层都会将浮点输入强制转换为其计算 dtype。

检索模型的初始权重。这样可以通过加载权重来从头开始训练。

initial_weights = model.get_weights()

使用 Model.fit 训练模型

接下来训练模型。

history = model.fit(x_train, y_train,
                    batch_size=8192,
                    epochs=5,
                    validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
print('Test loss:', test_scores[0])
print('Test accuracy:', test_scores[1])

Epoch 1/5
6/6 [==============================] - 0s 55ms/step - loss: 4.9447 - accuracy: 0.4281 - val_loss: 0.8652 - val_accuracy: 0.7736
Epoch 2/5
6/6 [==============================] - 0s 26ms/step - loss: 0.6878 - accuracy: 0.7959 - val_loss: 0.3972 - val_accuracy: 0.8672
Epoch 3/5
6/6 [==============================] - 0s 26ms/step - loss: 0.3874 - accuracy: 0.8704 - val_loss: 0.2237 - val_accuracy: 0.9322
Epoch 4/5
6/6 [==============================] - 0s 26ms/step - loss: 0.3237 - accuracy: 0.9021 - val_loss: 0.3711 - val_accuracy: 0.8843
Epoch 5/5
6/6 [==============================] - 0s 25ms/step - loss: 0.2369 - accuracy: 0.9287 - val_loss: 0.2498 - val_accuracy: 0.9153
313/313 - 1s - loss: 0.2635 - accuracy: 0.9079
Test loss: 0.26349955797195435
Test accuracy: 0.9078999757766724

请注意,模型会在日志中打印每个样本的时间:例如,“4us/sample”。第一个周期 (Epoch) 可能比较慢,因为 TensorFlow 需要花一些时间来优化模型,但之后每个样本的时间会稳定下来。

如果在 Colab 中运行本指南中的示例,您可以使用 float32 对比混合精度的性能。为此,请在“Setting the dtype policy”部分将策略从 mixed_float16 更改为 float32,然后重新运行所有代码单元,直到此代码点。在计算能力至少为 7.0 的 GPU 上,您会发现每个样本的时间大大增加,表明混合精度提升了模型的速度。例如,对于 Titan V GPU,每样本时间可从 4us 增加到 12us。在继续学习本指南之前,请确保将策略改回 mixed_float16 并重新运行代码单元。

对于很多实际模型,使用混合精度时还可以将批次大小加倍而不会耗尽内存,因为 float16 张量只需要使用 float32 一半的内存。不过,这对本文中所讲的小模型毫无意义,因为您几乎可以使用任何 dtype 来运行该模型,而每个批次可以包含有 60,000 张图片的整个 MNIST 数据集。

如果在 TPU 上运行混合精度,您会发现与在 GPU 上运行混合精度相比,性能提升并不明显。这是因为即使默认 dtype 策略为 float32,TPU 也会在后台执行一些 bfloat16 运算。对于某些使用 bfloat16 可以获得数值稳定性的运算(如矩阵乘法),TPU 硬件不支持使用 float32。对于此类运算,TPU 后端会自动在内部使用 bfloat16。因此,将 dtype='float32' 传递至使用此类运算的层不会产生数值影响,不过,使用 bfloat16 计算来运行此类层也不会产生负面影响。

损失放大

损失放大是 tf.keras.Model.fit 自动使用 mixed_float16 策略执行,从而避免数值下溢的一种技巧。本节介绍损失放大及其行为定制方法。

下溢和溢出

float16 数据类型的动态范围比 float32 窄。这意味着大于 $65504$ 的数值会因溢出而变为无穷大,小于 $6.0 \times 10^{-8}$ 的数值则会因下溢而变成零。float32 和 bfloat16 的动态范围要大得多,因此一般不会出现下溢或溢出的问题。

例如:

x = tf.constant(256, dtype='float16')
(x ** 2).numpy()  # Overflow
inf
x = tf.constant(1e-5, dtype='float16')
(x ** 2).numpy()  # Underflow
0.0

实际上,float16 也极少出现下溢的情况。此外,在正向传递中出现下溢的情形更是十分罕见。但是,在反向传递中,梯度可能因下溢而变为零。损失放大就是一个防止出现下溢的技巧。

损失放大背景知识

损失放大的基本概念非常简单:只需将损失乘以某个大数字即可,如 $1024$。我们将该数字称为损失标度。这会将梯度放大 $1024$ 倍,大大降低了发生下溢的几率。计算出最终梯度后,将其除以 $1024$ 即可得到正确值。

该过程的伪代码是:

loss_scale = 1024 loss = model(inputs) loss *= loss_scale # We assume `grads` are float32. We do not want to divide float16 gradients grads = compute_gradient(loss, model.trainable_variables) grads /= loss_scale

选择合适的损失标度比较困难。如果损失标度太小,梯度可能仍会因下溢而变为零。如果太大,则会出现相反的问题:梯度可能因溢出而变为无穷大。

为了解决这一问题,TensorFlow 会动态确定损失标度,因此,您不必人工选择。如果使用 tf.keras.Model.fit,则会自动完成损失放大,您不必做任何额外的工作。下一节会对此进行详细阐述。

选择损失标度

每一条 dtype 策略都有一个关联的 tf.mixed_precision.experimental.LossScale 对象(可选),这代表一个固定或动态损失标度。默认情况下,mixed_float16 策略的损失标度是 tf.mixed_precision.experimental.DynamicLossScale,它会动态确定损失标度值。其他策略在默认情况下没有损失标度,因为只有在使用 float16 时才需要该值。您可以查询策略的损失标度:

loss_scale = policy.loss_scale
print('Loss scale: %s' % loss_scale)
Loss scale: DynamicLossScale(current_loss_scale=32768.0, num_good_steps=30, initial_loss_scale=32768.0, increment_period=2000, multiplier=2.0)

损失标度会打印很多内部状态信息,不过您可以忽略。最重要的部分是 current_loss_scale,这里显示的就是损失标度的当前值。

在构建 dtype 策略时,您也可以通过传递一个数字来使用静态损失标度。

new_policy = mixed_precision.Policy('mixed_float16', loss_scale=1024)
print(new_policy.loss_scale)
FixedLossScale(1024.0)

dtype 策略的构造函数始终会将损失标度转换为 LossScale 对象。在本例中,它会转换为 tf.mixed_precision.experimental.FixedLossScale,即除 DynamicLossScale 之外的唯一 LossScale 子类。

注:建议使用动态损失标度,而不要使用任何其他值。选择固定损失标度非常困难,因为太小则会导致模型无法得到训练,而太大又会导致梯度中出现 Inf(无穷大)或 NaN(无效)值。动态损失标度的值通常非常接近最佳值,因此,您不必做任何修改。目前,动态损失标度的计算速度略低于固定损失标度,但未来会有所改善。

就像层一样,每个模型都有一条 dtype 策略。如果存在,模型会使用其策略的损失标度,在 tf.keras.Model.fit 方法中应用损失放大。这意味着如果使用 Model.fit,则完全不必担心损失放大:默认情况下,mixed_float16 策略有一个动态损失标度,而 Model.fit 将应用该值。

对于自定义训练循环,模型会忽略该策略的损失标度,因此,您必须手动应用该值。下一节会对此进行阐述。

使用自定义训练循环训练模型

目前,您已经使用混合精度通过 tf.keras.Model.fit 对 Keras 模型进行了训练。接下来将使用混合精度执行自定义训练循环。如果您还不了解自定义训练循环,请先阅读自定义训练指南

使用混合精度运行自定义训练循环需要对使用 float32 运行训练的模型进行两方面的更改:

  1. 使用混合精度构建模型(已完成)
  2. 如果使用 mixed_float16,则明确使用损失放大。

对于步骤 (2),您将使用 tf.keras.mixed_precision.experimental.LossScaleOptimizer 类,其中会封装一个优化器并应用损失放大。该类有两个参数:优化器和损失标度。像下面这样构造一个类即可使用动态损失标度

optimizer = keras.optimizers.RMSprop()
optimizer = mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')

传递 'dynamic' 等效于传递 tf.mixed_precision.experimental.DynamicLossScale()

接下来定义损失对象和 tf.data.Dataset

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
train_dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
                 .shuffle(10000).batch(8192))
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(8192)

接下来定义训练步骤函数。为了放大损失和缩小梯度,将使用损失标度优化器的两个新方法:

  • get_scaled_loss(loss):将损失值乘以损失标度值
  • get_unscaled_gradients(gradients):获取一系列放大的梯度作为输入,并将每一个梯度除以损失标度,从而将其缩小为实际值

为了防止梯度发生下溢,必须使用这些函数。随后,如果全部没有出现 Inf 或 NaN 值,则 LossScaleOptimizer.apply_gradients 会应用这些梯度。它还会更新损失标度,如果梯度出现 Inf 或 NaN 值,则会将其减半,而如果出现零值,则会增大损失标度。

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

在训练的开始阶段,LossScaleOptimizer 可能会跳过前几个步骤。先使用非常大的损失标度,以便快速确定最佳值。经过几个步骤后,损失标度将稳定下来,这时跳过的步骤将会很少。这一过程是自动执行的,不会影响训练质量。

现在定义测试步骤。

@tf.function
def test_step(x):
  return model(x, training=False)

加载模型的初始权重,以便从头开始重新训练。

model.set_weights(initial_weights)

最后,运行自定义训练循环。

for epoch in range(5):
  epoch_loss_avg = tf.keras.metrics.Mean()
  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='test_accuracy')
  for x, y in train_dataset:
    loss = train_step(x, y)
    epoch_loss_avg(loss)
  for x, y in test_dataset:
    predictions = test_step(x)
    test_accuracy.update_state(y, predictions)
  print('Epoch {}: loss={}, test accuracy={}'.format(epoch, epoch_loss_avg.result(), test_accuracy.result()))
Epoch 0: loss=3.7786550521850586, test accuracy=0.8511999845504761
Epoch 1: loss=0.502143919467926, test accuracy=0.920199990272522
Epoch 2: loss=0.30543556809425354, test accuracy=0.8456000089645386
Epoch 3: loss=0.296634316444397, test accuracy=0.945900022983551
Epoch 4: loss=0.2582075893878937, test accuracy=0.9528999924659729

GPU 性能提示

下面是在 GPU 上使用混合精度时的一些性能提示。

增大批次大小

当使用混合精度时,如果不影响模型质量,可以尝试使用双倍批次大小运行。因为 float16 张量只使用一半内存,所以,您通常可以将批次大小增大一倍,而不会耗尽内存。增大批次大小通常可以提高训练吞吐量,即模型每秒可以运行的训练元素数量。

确保使用 GPU Tensor 核心

如前所述,现代 NVIDIA GPU 使用称为 Tensor 核心的特殊硬件单元, 可以非常快速地执行 float16 矩阵乘法运算。但是,Tensor 核心要求张量的某些维度是 8 的倍数。在下面的示例中,当且仅当参数值是 8 的倍数时,才能使用 Tensor 核心。

  • tf.keras.layers.Dense(units=64)
  • tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)
    • 其他卷积层也是如此,如 tf.keras.layers.Conv3d
  • tf.keras.layers.LSTM(units=64)
    • 其他 RNN 也是如此,如 tf.keras.layers.GRU
  • tf.keras.Model.fit(epochs=2, batch_size=128)

您应该尽可能使用 Tensor 核心。如果要了解更多信息,请参阅 NVIDIA 深度学习性能指南,其中介绍了使用 Tensor 核心的具体要求以及与 Tensor 核心相关的其他性能信息。

XLA

XLA 是一款可以进一步提高混合精度性能,也可以在较小程度上提高 float32 性能的编译器。有关详细信息,请参阅 XLA 指南

Cloud TPU 性能提示

就像在 GPU 上一样,您也应该尝试将批次大小增大一倍,因为 bfloat16 张量同样只使用一半内存。双倍批次大小可能会提高训练吞吐量。

TPU 不需要任何其他特定于混合精度的调整即可获得最佳性能。TPU 已经要求使用 XLA。它们可以从某些是 $128$ 的倍数的维度获得优势,不过就像使用混合精度一样,这同样适用于 float32。有关一般 TPU 性能提示,请参阅Cloud TPU 性能指南。这些提示对混合精度和 float32 均适用。

总结

  • 如果您使用的是计算能力至少为 7.0 的 TPU 或 NVIDIA GPU,则应使用混合精度,因为它可以将性能提升多达 3 倍。
  • 您可以按如下代码使用混合精度:

    # On TPUs, use 'mixed_bfloat16' instead policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16') mixed_precision.set_policy(policy)
    
  • 如果您的模型以 softmax 结尾,请确保其类型为 float32。不管您的模型以什么结尾,必须确保输出为 float32。

  • 如果您通过 mixed_float16 使用自定义训练循环,则除了上述几行代码外,您还需要使用 tf.keras.mixed_precision.experimental.LossScaleOptimizer 封装您的优化器。然后调用 optimizer.get_scaled_loss 来放大损失,并且调用 optimizer.get_unscaled_gradients 来缩小梯度。

  • 如果不会降低计算准确率,则可以将训练批次大小加倍。

  • 在 GPU 上,确保大部分张量维度是 $8$ 的倍数,从而最大限度提高性能

有关使用 tf.keras.mixed_precision API 的混合精度的更多示例,请参阅官方模型仓库。大多数官方模型(如 ResNetTransformer)通过传递 --dtype=fp16 来使用混合精度。