混合精度

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 は 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

2022-12-14 22:52:34.442474: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 22:52:34.442569: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 22:52:34.442579: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

対応ハードウェア

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

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

注意: このガイドを Google Colab で実行する場合、通常、GPU ランタイムには P100 が接続されています。 P100 のコンピューティング機能は 6.0 なので、大幅なスピードアップは期待できません。

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

nvidia-smi -L

GPU 0: Tesla P100-PCIE-16GB (UUID: GPU-801d62f9-f9a1-dff7-35aa-53e4a0925e05)
GPU 1: Tesla P100-PCIE-16GB (UUID: GPU-97416b7b-22fd-9b90-b78e-89cd89f95f7e)
GPU 2: Tesla P100-PCIE-16GB (UUID: GPU-385f8592-5071-fd9b-8017-9fc41ae9cf87)
GPU 3: Tesla P100-PCIE-16GB (UUID: GPU-16291778-3dce-a194-61c3-016b4398130b)

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

高速化が期待できない CPU や古い GPU でも、単体テスト、デバッグ、または単に API を試す目的で、混合精度 API の使用は可能です。ただし CPU での混合精度の実行は、大幅に遅くなります。

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)

WARNING:tensorflow:Mixed precision compatibility check (mixed_float16): WARNING
Your GPUs may run slowly with dtype policy mixed_float16 because they do not have compute capability of at least 7.0. Your GPUs:
  Tesla P100-PCIE-16GB, compute capability 6.0 (x4)
See https://developer.nvidia.com/cuda-gpus for a list of GPUs and their compute capabilities.
If you will use compatible GPU(s) not attached to this host, e.g. by running a multi-worker model, you can ignore this warning. This message will only be logged once

つまり、文字列を直接 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 の場合は、代わりに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
6/6 [==============================] - 3s 203ms/step - loss: 2.4728 - accuracy: 0.3782 - val_loss: 0.9297 - val_accuracy: 0.7632
Epoch 2/5
6/6 [==============================] - 1s 157ms/step - loss: 1.0029 - accuracy: 0.7132 - val_loss: 0.4664 - val_accuracy: 0.8827
Epoch 3/5
6/6 [==============================] - 1s 156ms/step - loss: 0.5843 - accuracy: 0.8226 - val_loss: 0.5223 - val_accuracy: 0.8356
Epoch 4/5
6/6 [==============================] - 1s 156ms/step - loss: 0.3808 - accuracy: 0.8892 - val_loss: 0.4449 - val_accuracy: 0.8586
Epoch 5/5
6/6 [==============================] - 1s 156ms/step - loss: 0.4945 - accuracy: 0.8526 - val_loss: 0.2652 - val_accuracy: 0.9202
313/313 - 1s - loss: 0.2700 - accuracy: 0.9209 - 623ms/epoch - 2ms/step
Test loss: 0.26999878883361816
Test accuracy: 0.9208999872207642

モデルはステップあたりの時間をログに出力します（例:「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 でモデルを実行できるため、これはこのトイモデルには適用されません。

Loss Scaling（損失スケーリング）

損失スケーリングは、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
# 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.9650015830993652, test accuracy=0.5597000122070312
Epoch 1: loss=0.8010242581367493, test accuracy=0.7914000153541565
Epoch 2: loss=0.511243462562561, test accuracy=0.8539000153541565
Epoch 3: loss=0.352459579706192, test accuracy=0.9063000082969666
Epoch 4: loss=0.3114456236362457, test accuracy=0.9176999926567078

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 を使用している場合、混合精度を使用することでパフォーマンスを最大 3 倍向上できる。

  • 混合精度は次の行で使用できる。

    # On TPUs, 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 を呼び出して勾配をスケーリング解除する。

  • 評価精度が低下しないようであれば、トレーニングのバッチサイズを 2 倍にする。

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

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