混合精度

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

總覽

混合精度是在模型訓練期間同時使用 16 位元和 32 位元浮點類型，以加快執行速度並減少記憶體用量。藉由將模型的特定部分保留在 32 位元類型中以維持數值穩定性，模型將具有較低的步長時間，並在評估指標 (例如準確度) 方面訓練得同樣良好。本指南說明如何使用 Keras 混合精度 API 來加速您的模型。使用此 API 可將現代 GPU 的效能提升 3 倍以上、TPU 的效能提升 60%，以及最新 Intel CPU 的效能提升 2 倍以上。

如今，大多數模型都使用 float32 dtype，佔用 32 位元記憶體。但是，有兩種精確度較低的 dtype，float16 和 bfloat16，各自佔用 16 位元記憶體。現代加速器可以更快速地在 16 位元 dtype 中執行運算，因為它們具有專門的硬體來執行 16 位元計算，而且可以更快地從記憶體讀取 16 位元 dtype。

NVIDIA GPU 可以在 float16 中比在 float32 中更快速地執行運算，而 TPU 和支援的 Intel CPU 可以在 bfloat16 中比 float32 更快速地執行運算。因此，在這些裝置上應盡可能使用這些精確度較低的 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、Cloud TPU 和最新 Intel CPU 上加速模型。NVIDIA GPU 支援混合使用 float16 和 float32，而 TPU 和 Intel CPU 支援混合使用 bfloat16 和 float32。

在 NVIDIA GPU 中，運算能力為 7.0 或更高的 GPU 將從混合精度中獲得最大的效能優勢，因為它們具有稱為 Tensor Core 的特殊硬體單元，可加速 float16 矩陣乘法和卷積。較舊的 GPU 在使用混合精度方面不會提供數學效能優勢，但是記憶體和頻寬節省可以實現一些加速。您可以在 NVIDIA 的 CUDA GPU 網頁上查詢 GPU 的運算能力。將從混合精度中受益最多的 GPU 範例包括 RTX GPU、V100 和 A100。

在 Intel CPU 中，從第 4 代 Intel Xeon 處理器 (代號 Sapphire Rapids) 開始，將從混合精度中獲得最大的效能優勢，因為它們可以使用 AMX 指令加速 bfloat16 計算 (需要 Tensorflow 2.12 或更高版本)。

您可以使用以下命令檢查 GPU 類型。只有在安裝 NVIDIA 驅動程式時，該命令才存在，否則以下命令會引發錯誤。

nvidia-smi -L

所有 Cloud TPU 都支援 bfloat16。

即使在預計不會加速的較舊 Intel 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 執行。模型的第一層會將輸入轉換為 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])

請注意，模型會在記錄中列印每個步驟的時間：例如，「25 毫秒/步驟」。第一個 epoch 可能會較慢，因為 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 之前的 GPU。這是因為 TPU 即使在使用預設 dtype 政策 float32 的情況下，也會在幕後以 bfloat16 執行某些運算。這類似於 Ampere GPU 預設使用 TensorFloat-32 的方式。與 Ampere GPU 相比，TPU 通常在實際模型中使用混合精度時看到的效能提升較少。

對於許多實際模型，混合精度也允許您將批次大小加倍而不會耗盡記憶體，因為 float16 張量佔用一半的記憶體。但是，這不適用於此玩具模型，因為您可能會在任何 dtype 中執行模型，其中每個批次都包含整個 MNIST 資料集 (包含 60,000 張圖片)。

損失縮放

損失縮放是一種技術，tf.keras.Model.fit 會自動對 mixed_float16 政策執行損失縮放，以避免數值下溢。本節說明什麼是損失縮放，下一節說明如何在自訂訓練迴圈中使用它。

下溢和溢位

與 float32 相比，float16 資料類型具有較窄的動態範圍。這表示高於 \(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 執行自訂訓練迴圈進行兩項變更

  1. 使用混合精度建構模型 (您已經執行了此操作)
  2. 如果使用 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.Datasets

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 Tensor Core

如先前所述，現代 NVIDIA GPU 使用稱為 Tensor Core 的特殊硬體單元，可以非常快速地乘法 float16 矩陣。但是，Tensor Core 要求張量的特定維度是 8 的倍數。在以下範例中，只有在需要使用 Tensor Core 時，引數才會以粗體顯示。

  • 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 Core。如果您想瞭解更多資訊，NVIDIA 深度學習效能指南說明了使用 Tensor Core 的確切要求以及其他與 Tensor Core 相關的效能資訊。

XLA

XLA 是一種編譯器，可以進一步提高混合精度效能，以及在較小程度上提高 float32 效能。請參閱 XLA 指南以瞭解詳細資訊。

Cloud TPU 效能提示

與 GPU 相同，您應嘗試在使用 Cloud TPU 時將批次大小加倍，因為 bfloat16 張量使用一半的記憶體。將批次大小加倍可能會增加訓練輸送量。

TPU 不需要任何其他特定於混合精度的調整即可獲得最佳效能。它們已經需要使用 XLA。TPU 受益於某些維度是 \(128\) 的倍數，但這同樣適用於 float32 類型，也適用於混合精度。查看 Cloud TPU 效能指南以瞭解一般 TPU 效能提示，這些提示既適用於混合精度，也適用於 float32 張量。

摘要

  • 如果您使用 TPU、運算能力至少為 7.0 的 NVIDIA GPU，或支援 AMX 指令集的 Intel 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。