混合精度

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

概述

混合精度是在训练期间在模型中使用 16 位和 32 位浮点类型，以使其运行速度更快并使用更少的内存。通过将模型的某些部分保留在 32 位类型中以确保数值稳定性，模型将具有更低的步长时间，并且在评估指标（如准确率）方面训练效果相同。本指南介绍了如何使用 Keras 混合精度 API 来加速您的模型。使用此 API 可以将现代 GPU 上的性能提高 3 倍以上，TPU 上的性能提高 60%，最新英特尔 CPU 上的性能提高 2 倍以上。

如今，大多数模型使用 float32 dtype，它占用 32 位内存。但是，还有两种精度较低的 dtype，float16 和 bfloat16，它们分别占用 16 位内存。现代加速器可以在 16 位 dtype 中更快地运行操作，因为它们具有专门的硬件来运行 16 位计算，并且 16 位 dtype 可以更快地从内存中读取。

NVIDIA GPU 可以比 float32 更快地在 float16 中运行操作，而 TPU 和支持的英特尔 CPU 可以比 float32 更快地在 bfloat16 中运行操作。因此，应尽可能在这些设备上使用这些精度较低的 dtype。但是，出于数值原因，变量和一些计算仍应在 float32 中进行，以便模型训练到相同的质量。Keras 混合精度 API 允许您将 float16 或 bfloat16 与 float32 混合使用，以获得 float16/bfloat16 的性能优势和 float32 的数值稳定性优势。

设置

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import mixed_precision

支持的硬件

虽然混合精度可以在大多数硬件上运行，但它只会加速最新 NVIDIA GPU、云 TPU 和最新英特尔 CPU 上的模型。NVIDIA GPU 支持使用 float16 和 float32 的混合，而 TPU 和英特尔 CPU 支持使用 bfloat16 和 float32 的混合。

在 NVIDIA GPU 中，计算能力为 7.0 或更高的 GPU 将从混合精度中获得最大的性能提升，因为它们具有特殊的硬件单元（称为张量核心）来加速 float16 矩阵乘法和卷积。较旧的 GPU 在使用混合精度时不会提供任何数学性能优势，但内存和带宽节省可以实现一些加速。您可以在 NVIDIA 的 CUDA GPU 网页上查找您的 GPU 的计算能力。从混合精度中获益最大的 GPU 示例包括 RTX GPU、V100 和 A100。

在英特尔 CPU 中，从第 4 代英特尔至强处理器（代号为 Sapphire Rapids）开始，将从混合精度中获得最大的性能提升，因为它们可以使用 AMX 指令（需要 Tensorflow 2.12 或更高版本）加速 bfloat16 计算。

您可以使用以下方法检查您的 GPU 类型。该命令仅在安装了 NVIDIA 驱动程序的情况下才存在，因此如果未安装，以下操作将引发错误。

nvidia-smi -L

所有云 TPU 都支持 bfloat16。

即使在较旧的英特尔 CPU、没有 AMX 的其他 x86 CPU 和较旧的 GPU 上，预计不会出现加速，混合精度 API 仍然可以用于单元测试、调试或只是尝试 API。但是，在没有 AMX 指令的 CPU 上使用 mixed_bfloat16 以及在所有 x86 CPU 上使用 mixed_float16 将运行得明显更慢。

设置 dtype 策略

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

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

简而言之，您可以直接将字符串传递给 set_global_policy，这通常在实践中完成。

# Equivalent to the two lines above
mixed_precision.set_global_policy('mixed_float16')

该策略指定了层面的两个重要方面：层计算所使用的 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)

如前所述，mixed_float16 策略将最显着地提高计算能力至少为 7.0 的 NVIDIA GPU 的性能。该策略将在其他 GPU 和 CPU 上运行，但可能不会提高性能。对于 TPU 和 CPU，应使用 mixed_bfloat16 策略。

构建模型

接下来，让我们开始构建一个简单的模型。非常小的玩具模型通常不会从混合精度中受益，因为 TensorFlow 运行时的开销通常会主导执行时间，从而使 GPU 上的任何性能改进都微不足道。因此，如果使用 GPU，让我们构建两个具有 4096 个单元的大型 Dense 层。

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)

每个层都有一个策略，默认情况下使用全局策略。因此，每个 Dense 层都具有 mixed_float16 策略，因为您之前将全局策略设置为 mixed_float16。这将导致密集层进行 float16 计算并具有 float32 变量。它们将输入转换为 float16 以进行 float16 计算，这会导致它们的输出也为 float16。它们的变量是 float32，在调用层时将转换为 float16，以避免 dtype 不匹配导致的错误。

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

接下来，创建输出预测。通常，您可以按如下方式创建输出预测，但这在 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)

模型末尾的 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)

将 dtype='float32' 传递给 softmax 层构造函数会覆盖层的 dtype 策略，使其成为 float32 策略，该策略在 float32 中进行计算并保留变量。等效地，您也可以传递 dtype=mixed_precision.Policy('float32')；层始终将 dtype 参数转换为策略。由于 Activation 层没有变量，因此策略的变量 dtype 被忽略，但策略的 float32 计算 dtype 会导致 softmax 和模型输出为 float32。

在模型中间添加 float16 softmax 是可以的，但在模型末尾的 softmax 应为 float32。原因是，如果从 softmax 流向损失的中间张量是 float16 或 bfloat16，则可能会出现数值问题。

如果您认为 float16 计算在数值上不稳定，可以通过传递 dtype='float32' 来覆盖任何层的 dtype 为 float32。但通常，这仅在模型的最后一层是必要的，因为大多数层在 mixed_float16 和 mixed_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，因为除以 255 的操作是在 CPU 上执行的，而 CPU 上的 float16 操作比 float32 操作运行得更慢。在这种情况下，性能差异可以忽略不计，但总的来说，如果输入处理数学运算在 CPU 上运行，则应在 float32 中运行。

检索模型的初始权重。这将允许通过加载权重从头开始训练。

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])

请注意，模型在日志中打印了每步的时间：例如，“25ms/step”。第一个纪元可能比较慢，因为 TensorFlow 会花费一些时间优化模型，但之后每步的时间应该稳定下来。

如果您在 Colab 中运行本指南，则可以比较混合精度与 float32 的性能。为此，将“设置 dtype 策略”部分中的策略从 mixed_float16 更改为 float32，然后重新运行所有单元格，直到此点。在计算能力为 7.X 的 GPU 上，您应该看到每步的时间显着增加，这表明混合精度加速了模型。在继续本指南之前，请确保将策略更改回 mixed_float16 并重新运行单元格。

在计算能力至少为 8.0 的 GPU（Ampere GPU 及更高版本）上，当使用混合精度与 float32 相比时，您可能不会在本指南中的玩具模型中看到性能改进。这是由于使用了 TensorFloat-32，它会在某些 float32 操作（如 tf.linalg.matmul）中自动使用较低的精度数学运算。TensorFloat-32 在使用 float32 时提供了一些混合精度的性能优势。但是，在现实世界的模型中，您仍然会通常从混合精度中体验到显着的性能改进，这是由于内存带宽节省和 TensorFloat-32 不支持的操作。

如果在 TPU 上运行混合精度，与在 GPU 上运行混合精度相比，您不会看到那么大的性能提升，尤其是对于 Ampere 之前的 GPU。这是因为 TPU 在后台以 bfloat16 执行某些操作，即使默认 dtype 策略为 float32 也是如此。这类似于 Ampere GPU 默认使用 TensorFloat-32 的方式。与 Ampere GPU 相比，TPU 在现实世界的模型上使用混合精度通常会看到更少的性能提升。

对于许多现实世界的模型，混合精度还允许您将批次大小加倍而不会出现内存不足的情况，因为 float16 张量占用的内存只有一半。但是，这并不适用于此玩具模型，因为您可能可以在任何 dtype 中运行模型，其中每个批次包含整个 MNIST 数据集（60,000 张图像）。

损失缩放

损失缩放是一种技术，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

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

在实践中，float16 的上溢很少发生。此外，下溢在正向传播期间也很少发生。但是，在反向传播期间，梯度可能会下溢到零。损失缩放是一种防止这种下溢的技术。

损失缩放概述

损失缩放的基本概念很简单：只需将损失乘以某个大数，例如 \(1024\)，您就会得到损失缩放值。这将导致梯度也按 \(1024\) 缩放，从而大大降低下溢的可能性。一旦计算出最终梯度，就将它们除以 \(1024\) 以将它们恢复到正确的值。

此过程的伪代码如下

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

选择损失缩放值可能很棘手。如果损失缩放值太低，梯度可能仍然下溢到零。如果太高，就会出现相反的问题：梯度可能会上溢到无穷大。

为了解决这个问题，TensorFlow 会动态确定损失缩放值，因此您无需手动选择。如果您使用 tf.keras.Model.fit，损失缩放会为您完成，因此您无需执行任何额外操作。如果您使用自定义训练循环，则必须显式使用特殊优化器包装器 tf.keras.mixed_precision.LossScaleOptimizer 才能使用损失缩放。这将在下一节中介绍。

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

到目前为止，您已经使用 tf.keras.Model.fit 使用混合精度训练了 Keras 模型。接下来，您将使用混合精度和自定义训练循环。如果您还不了解自定义训练循环是什么，请先阅读自定义训练指南。

使用混合精度运行自定义训练循环需要对在 float32 中运行它进行两个更改

使用混合精度构建模型（您已经完成了此操作）
如果使用 mixed_float16，则显式使用损失缩放。

对于步骤 (2)，您将使用 tf.keras.mixed_precision.LossScaleOptimizer 类，它包装了一个优化器并应用损失缩放。默认情况下，它会动态确定损失缩放值，因此您无需选择。按如下方式构造 LossScaleOptimizer。

optimizer = keras.optimizers.RMSprop()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)

如果您愿意，可以选择显式损失缩放值或以其他方式自定义损失缩放行为，但强烈建议保留默认损失缩放行为，因为它已被发现适用于所有已知模型。如果您想自定义损失缩放行为，请参阅 tf.keras.mixed_precision.LossScaleOptimizer 文档。

接下来，定义损失对象和 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): 接收一个缩放梯度列表作为输入，并将每个梯度除以损失尺度以取消缩放。

必须使用这些函数来防止梯度下溢。 LossScaleOptimizer.apply_gradients 然后将应用梯度，前提是它们都不包含 Inf 或 NaN。它还会更新损失尺度，如果梯度包含 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()))

GPU 性能提示

以下是使用混合精度在 GPU 上进行训练时的性能提示。

增加批次大小

如果不会影响模型质量，请尝试在使用混合精度时将批次大小增加一倍。由于 float16 张量使用一半的内存，这通常允许您将批次大小增加一倍而不会出现内存不足的情况。增加批次大小通常会提高训练吞吐量，即模型每秒可以运行的训练元素数量。

确保使用 GPU 张量核心

如前所述，现代 NVIDIA GPU 使用称为张量核心的特殊硬件单元，可以非常快速地乘以 float16 矩阵。但是，张量核心要求张量的某些维度必须是 8 的倍数。在以下示例中，如果参数需要是 8 的倍数才能使用张量核心，则该参数将以粗体显示。

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)

您应该尽可能地使用张量核心。如果您想了解更多信息，NVIDIA 深度学习性能指南描述了使用张量核心的确切要求以及其他与张量核心相关的性能信息。

XLA

XLA 是一个编译器，可以进一步提高混合精度性能，以及 float32 性能（程度较小）。有关详细信息，请参阅 XLA 指南。

Cloud TPU 性能提示

与 GPU 一样，您应该尝试在使用 Cloud TPU 时将批次大小增加一倍，因为 bfloat16 张量使用一半的内存。增加批次大小可能会提高训练吞吐量。

TPU 不需要任何其他与混合精度相关的调整即可获得最佳性能。它们已经要求使用 XLA。TPU 从某些维度是 \(128\) 的倍数中受益，但这同样适用于 float32 类型，也适用于混合精度。查看 Cloud TPU 性能指南，了解适用于混合精度和 float32 张量的通用 TPU 性能提示。

总结

如果您使用 TPU、具有至少 7.0 计算能力的 NVIDIA GPU 或支持 AMX 指令的英特尔 CPU，则应该使用混合精度，因为它可以将性能提高多达 3 倍。

您可以使用以下几行代码使用混合精度

# On TPUs and CPUs, use 'mixed_bfloat16' instead
mixed_precision.set_global_policy('mixed_float16')

如果您的模型以 softmax 结束，请确保它是 float32。无论您的模型以什么结束，请确保输出是 float32。
如果您使用带有 mixed_float16 的自定义训练循环，除了上面的几行代码之外，您还需要将优化器包装在 tf.keras.mixed_precision.LossScaleOptimizer 中。然后调用 optimizer.get_scaled_loss 来缩放损失，并调用 optimizer.get_unscaled_gradients 来取消缩放梯度。
如果您使用带有 mixed_bfloat16 的自定义训练循环，设置上面提到的 global_policy 就足够了。
如果不会降低评估精度，请将训练批次大小增加一倍
在 GPU 上，确保大多数张量维度是 \(8\) 的倍数，以最大限度地提高性能

有关使用 tf.keras.mixed_precision API 的混合精度的示例，请查看与训练性能相关的函数和类。查看官方模型，例如 Transformer，了解详细信息。