混合精度

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード

概要

混合精度とは、トレーニング中に 16 ビットと 32 ビット浮動小数点型の両方をモデルに使ってモデルのトレーニングを高速化し、使用するメモリを少なくする手法です。数値の安定性を保つためにモデルの特定の部分を 32 ビットで保持することにより、モデルのステップ時間を短縮しつつ、精度などの評価指標は同様にトレーニングすることができます。このガイドでは、Keras 混合精度 API でモデルを高速化する使い方を説明します。この API を使用すると、最新の GPU で 3 倍以上、TPU で 60% 以上パフォーマンスを向上させることができます。

現在、殆どのモデルでは 32 ビットのメモリを必要とする float32 dtype が使用されています。しかしながら、代わりに必要とするメモリが 16 ビットの float16 と bfloat16 という 2 つの低精度 dtype があります。最近のアクセラレータは、16 ビットの計算実行専門のハードウェアを備えており、16 ビットの dtype はより高速でメモリから読み取ることができるため、 16 ビットの dtype 演算をより高速に実行できます。

NVIDIA GPU は float32 よりも float16 で速く演算を実行でき、TPU とサポートしている Intel CPU は float32 よりも bfloat16 で速く演算を実行できます。したがって、これらのデバイスでは低精度の dtype を可能な限り使用することが推奨されます。ただし、変数および一部の計算は、モデルのトレーニングの品質を維持するために、数値的理由から float32 のままにする必要があります。Keras 混合精度 API を使用すると、float16 または bfloat16 と float32 の組み合わせが可能になり、float16 / bfloat16 によるパフォーマンスのメリットと float32 による数値的安定性のメリットの両方を得ることができます。

注意: このガイドでは、「数値的安定性」という用語は、高精度の dtype ではなく低精度の dtype を使用することによって、モデルの品質がどのように影響を受けるかを指します。これらの dtype のいずれかで実行し、モデルの評価精度やその他のメトリクスが float32 での演算実行と比較して低下する場合、その演算は float16 または bfloat16 で「数値的に不安定」と言えます。

セットアップ

import tensorflow as tf

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

2024-01-11 18:54:25.810615: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-01-11 18:54:25.810663: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-01-11 18:54:25.812179: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

対応ハードウェア

混合精度はほとんどのハードウェアで動作しますが、高速化できるのは最近の NVIDIA GPU、Cloud TPU、または最近の Intel CPU のモデルに限ります。NVIDIA GPU は float16 と float32 を組み合わせた使用に、TPU と Intel CPU は bfloat16 と float32 を組み合わせた使用に対応しています。

NVIDIA GPUの中でも、コンピューティング機能が 7.0 以上の場合、float16 の行列乗算と畳み込みを加速するテンソルコアと呼ばれる特別なハードウェアユニットを備えているため、混合精度のパフォーマンスが最大になります。古い GPU の場合、混合精度の使用による数学的パフォーマンスは期待できませんが、メモリと帯域幅の節約によって幾分かの高速化は可能です。NVIDIA の CUDA GPU ウェブページから、お持ちの GPU のコンピューティング機能を確認できます。混合精度が最も有効な GPU には、RTX GPU、V100、A100 などがあります。

Intel CPU の間では、第 4 世代 Intel Xeon プロセッサ（コード名「Sapphire Rapids」）に混合精度による最大のパフォーマンスメリットが見られます。これらは、AMX 命令を使って bfloat16 の計算を高速化できるためです（Tensorflow 2.12 以降が必要です）。

注意: このガイドを Google Colab で実行する場合、GPU ランタイムには通常 P100 が接続されています。P100 のコンピューティング能力は 6.0 であるため大幅な高速化は期待されません。CPU ランタイムで実行する場合、このランタイムには AMX を使用しない CPU がある可能性があるため、低速化が発生する可能性があります。

GPU タイプは以下の方法で確認できます。このコマンドは NVIDIA ドライバがインストールされている場合にのみ存在するため、そうでない場合はエラーが生成されます。

nvidia-smi -L
 
GPU 0: Tesla T4 (UUID: GPU-a5cd732a-9281-3439-08ab-4a1133335283)
GPU 1: Tesla T4 (UUID: GPU-25bf817b-a042-aca4-1e88-497783f147e8)
GPU 2: Tesla T4 (UUID: GPU-c1f518b2-13c1-3c07-13e9-0e05033f7c53)
GPU 3: Tesla T4 (UUID: GPU-5e6dae29-da98-3180-fb9f-4654d5d3b4a6)

すべての Cloud TPU は bfloat16 に対応しています。

高速化が期待できない CPU、AMX を使用しないその他の x86 CPU、および古い GPU でも、単体テスト、デバッグ、または単に API を試す目的で、混合精度 API の使用は可能です。ただし AMX 命令なしの CPU での mixed_bfloat16 とすべての x86 CPU での mixed_bfloat16 の実行は、大幅に遅くなります。

dtype ポリシーを設定する

Keras で混合精度を使用するには、通常 dtype ポリシーと呼ばれるtf.keras.mixed_precision.Policy を作成する必要があります。dtype ポリシーは、実行される dtypes レイヤーを指定します。このガイドでは、文字列 'mixed_float16' からポリシーを作成し、それをグローバルポリシーとして設定します。これにより、その後に作成されるレイヤーは、float16 と float32 を組み合わせた混合精度を使用します。

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

INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0

つまり、文字列を直接 set_global_policy に渡すことができ、この方法は一般的に実践で行われます。

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

ポリシーは、レイヤーの計算が行われる dtype と、レイヤーの変数の dtype という、レイヤーの 2 つの重要な側面を指定します。上記では、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

前述したように、mixed_float16 ポリシーは、7.0 以上のコンピューティング能力を備えた NVIDIA GPU のパフォーマンスを最も大幅に向上させます。ポリシーは他の GPU や CPU でも実行できますが、パフォーマンスは向上しない可能性があります。TPU と CPU の場合は、代わりに mixed_bfloat16 ポリシーを使用する必要があります。

モデルを作成する

次に、簡単なモデルを作成してみましょう。非常に小さなトイモデルでは、通常 TensorFlow ランタイムのオーバーヘッドが実行時間を支配し、GPU のパフォーマンス向上がごく僅かになってしまうため、混合精度の恩恵を受けることができません。したがって、GPU を使用する場合には、それぞれが 4096 ユニットの 2 つの大きな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)

The model will run with 4096 units on a GPU

各レイヤーにはポリシーがあり、グローバルポリシーをデフォルトで使用します。前にグローバルポリシーをmixed_float16と設定したため、各Denseレイヤーにはmixed_float16ポリシーがあります。これによって、密なレイヤーは float16 計算を行い、float32 変数を持ちます。これらはfloat16 計算を行うために入力を float16 にキャストし、その結果、出力は float16 になります。変数は float32 なので、dtype の不一致によるエラーを回避するためにレイヤーを呼び出す際に、 float16 にキャストされます。

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)

<Policy "mixed_float16">
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であるため、ソフトマックスのアクティブ化は通常 float16 計算 dtype を持ち、float16 テンソルを出力します。

これは、高密度（Dense）レイヤーとソフトマックスレイヤーを分離して、ソフトマックスレイヤーにdtype='float32'を渡して修正することができます。

# 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'をソフトマックス レイヤー コンストラクタ―に渡すと、レイヤーの dtype ポリシーがfloat32ポリシーにオーバーライドされ、計算を行い、変数を float32 で保持します。同様に、その代わりにレイヤーが常に dtype 引数をポリシーに変換するdtype=mixed_precision.Policy('float32')を渡すこともできます。Activationレイヤーには変数がないため、ポリシーの変数 dtype は無視されますが、ポリシーの計算 dtype float32 はソフトマックスを適用し、モデル出力は float32 になります。

モデルの中間で float16 ソフトマックスを追加することは問題ありませんが、モデルの最後のソフトマックスは float32 にする必要があります。 その理由は、ソフトマックスから損失に流れる中間テンソルが float16 または bfloat16 である場合、数値の問題が発生する可能性があるためです。

float16 計算で数値的に安定しないと思われる場合、dtype='float32'を渡して任意のレイヤーの dtype を float32 にオーバーライドできます。ただし、ほとんどのレイヤーは mixed_float16 または mixed_bfloat16 で十分な精度があるため、通常は、モデルの最後のレイヤーのみで必要です。

モデルがソフトマックスで終わらない場合でも、出力は 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 にキャストします。255 による除算は CPU で行われ、CPU の float16 演算は float32 演算よりも実行速度が遅いため、float16 にはキャストしません。この場合のパフォーマンスの違いはごくわずかですが、一般に CPU で実行する場合には float32 で入力処理演算を実行する必要があります。各レイヤーが浮動小数点入力を dtype 計算にキャストするため、モデルの最初のレイヤーは、入力を float16 にキャストします。

モデルの重みの初期値が取得されます。これによって、重みをロードして最初からトレーニングを再開できます。

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
1/6 [====>.........................] - ETA: 7s - loss: 2.3131 - accuracy: 0.0955
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1704999272.626079  116647 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
6/6 [==============================] - 2s 133ms/step - loss: 2.3236 - accuracy: 0.3688 - val_loss: 0.9058 - val_accuracy: 0.7247
Epoch 2/5
6/6 [==============================] - 0s 65ms/step - loss: 1.0910 - accuracy: 0.6608 - val_loss: 0.4745 - val_accuracy: 0.8775
Epoch 3/5
6/6 [==============================] - 0s 65ms/step - loss: 0.5150 - accuracy: 0.8365 - val_loss: 0.4096 - val_accuracy: 0.8535
Epoch 4/5
6/6 [==============================] - 0s 66ms/step - loss: 0.4104 - accuracy: 0.8651 - val_loss: 0.4787 - val_accuracy: 0.8520
Epoch 5/5
6/6 [==============================] - 0s 65ms/step - loss: 0.3500 - accuracy: 0.8956 - val_loss: 0.2303 - val_accuracy: 0.9357
313/313 - 1s - loss: 0.2358 - accuracy: 0.9315 - 661ms/epoch - 2ms/step
Test loss: 0.23582278192043304
Test accuracy: 0.9315000176429749

モデルはステップあたりの時間をログに出力します（例:「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 で実行する場合に比べそれほどパフォーマンスのゲインは見られません。これは、TPU が、デフォルトの dtype ポリシーが float32 であっても内部的には bfloat16 で特定の演算を行うためです。これは Ampere GPU が TensorFloat-32 をデフォルトで使用する方法に似ています。Ampere GPU に比べ、TPU では通常は、実世界モデルの混合精度でパフォーマンスのゲインをあまり得られません。

float16 テンソルの使用メモリは半分で済むため、実世界の多くのモデルでは、バッチサイズを 2 倍にしてもメモリ不足にならずに混合精度の使用が可能です。ただし、60,000 枚の画像から成る MNIST データセット全体で構成されるバッチは任意の dtype でモデルを実行できるため、これはこのトイモデルには適用されません。

損失スケーリング

損失スケーリングは、tf.keras.Model.fit が mixed_float16 ポリシーを使用して自動的に実行し、数値のアンダーフローを回避する手法です。このセクションでは、損失スケーリングをカスタムトレーニングループと使用する方法について説明します。

注意: mixed_bfloat16 ポリシーを使用する場合、損失スケーリングを行う必要はありません。

アンダーフローとオーバーフロー

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
# 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 のみで実行する場合に比べ、2 つの変更が必要です。

  1. 混合精度でモデルを構築する（既に構築済み）
  2. mixed_float16が使用されている場合は、明示的に損失スケーリングを使用する

手順 (2) では、tf.keras.mixed_precision.LossScaleOptimizer クラスを使用し、オプティマイザをラップして損失スケーリングを適用します。デフォルトでは、損失スケールが動的に決定されるようになっているため、何も選択する必要はありません。次のようにして、これにはオプティマイザと損失スケールの 2 つの引数が必要です。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)

次に、トレーニングステップ関数を定義します。 損失をスケーリングし、勾配をスケーリング解除するために、損失スケールオプティマイザの 2 つの新しいメソッドを使用します。

  • 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=1.8230167627334595, test accuracy=0.5349000096321106
Epoch 1: loss=0.8649435639381409, test accuracy=0.7617999911308289
Epoch 2: loss=0.43499863147735596, test accuracy=0.8953999876976013
Epoch 3: loss=0.42511460185050964, test accuracy=0.9345999956130981
Epoch 4: loss=0.25523805618286133, test accuracy=0.8705999851226807

GPU パフォーマンスに関するヒント

GPU で混合精度を使用する場合のパフォーマンスに関するヒントをいくつか紹介します。

バッチサイズを大きくする

モデルの品質に影響がない場合は、混合精度の使用時にバッチサイズを 2 倍にして実行してみてください。float16 テンソルは半分のメモリを使用するため、これにより多くの場合はメモリを使い果たすことなくバッチサイズを 2 倍にすることができます。バッチサイズを増やすと、通常はトレーニングスループット、すなわちモデルで実行できる 1 秒あたりのトレーニング要素が増加します。

GPU のテンソルコアを使用する

前述したように、最近の NVIDIA GPU は、float16 行列を非常に速く乗算できる Tensor Core と呼ばれる特殊なハードウェアユニットを使用しています。ただしテンソルコアは、テンソルの特定の次元を 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)
    • tf.keras.layers.GRU など、他の RNN についても同様
  • tf.keras.Model.fit(epochs=2, batch_size=128)

可能な場合はできる限りテンソルコアを使用するようにしてください。詳細については、NVIDIA ディープラーニングパフォーマンスガイド（英語）において、テンソルコア使用のための正確な要件、その他テンソルコア関連のパフォーマンス情報などが記載されています。

XLA

XLA は、混合精度のパフォーマンスをさらに向上させるだけでなく、比較的規模は小さいですが float32 のパフォーマンスもさらに向上させることができるコンパイラです。詳しくは、XLA ガイドをご覧ください。

Cloud TPU パフォーマンスに関するヒント

GPU と同様に、bfloat16 テンソルは半分のメモリを使用するため、Cloud TPU を使用する場合はバッチサイズを 2 倍にしてみてください。バッチサイズを 2 倍にすると、トレーニングのスループットが向上する場合があります。

TPU は、パフォーマンスの最適化に他の混合精度固有の調整は必要ありません。TPU は既に XLA の使用が必須です。\(128\) の倍数という特定の次元を持っていることが有効で、これは混合精度の場合と同様に float32 にも同じように適用されます。混合精度と float32 テンソルに適用される一般的な TPU パフォーマンスに関するヒントについては、Cloud TPU パフォーマンスガイドをご覧ください。

まとめ

  • 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 を設定するだけで十分です。

  • Double the training batch size if it does not reduce evaluation accuracy

  • GPU では、パフォーマンスを最大化するためには、テンソルの次元の大部分を \(8\) の倍数にします。

tf.keras.mixed_precision API を使った混合精度の例については、トレーニングパフォーマンスに関連する関数とクラスをご覧ください。詳細は、Transformer などの公式モデルをご覧ください。