今日のローカルTensorFlowEverywhereイベントの出欠確認！

混合精度

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

概要

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

注意: Keras 混合精度 API は現在実験段階であり、変更される可能性があります。

現在、殆どのモデルでは 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 で「数値的に不安定」と言えます。

セットアップ

Keras 混合精度 API は TensorFlow 2.1 で利用可能です。

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 以上の場合、float16 行列の乗算と畳み込みを加速する Tensor Core と呼ばれる特別なハードウェアユニットを備えているため、混合精度のパフォーマンスが最大になります。古い 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-3a80b512-ac47-0257-1c4d-c6966adf0f87)

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

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

dtype ポリシーを設定する

Keras で混合精度を使用するには、通常 dtype ポリシーと呼ばれるtf.keras.mixed_precision.experimental.Policyを作成する必要があります。dtype ポリシーは、実行される dtypes レイヤーを指定します。このガイドでは、文字列'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
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/mixed_precision/loss_scale.py:56: DynamicLossScale.__init__ (from tensorflow.python.training.experimental.loss_scale) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.keras.mixed_precision.LossScaleOptimizer instead. LossScaleOptimizer now has all the functionality of DynamicLossScale

ポリシーは、レイヤーの計算が行われる 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('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であるため、ソフトマックスのアクティブ化は通常 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 にオーバーライドできます。ただし、ほとんどのレイヤーは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

WARNING:tensorflow:tf.keras.mixed_precision.experimental.LossScaleOptimizer is deprecated. Please use tf.keras.mixed_precision.LossScaleOptimizer instead. Note that the non-experimental LossScaleOptimizer does not take a DynamicLossScale but instead takes the dynamic configuration directly in the constructor. For example:
  opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt)

この例では、入力データを 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 [==============================] - 2s 113ms/step - loss: 5.0315 - accuracy: 0.3061 - val_loss: 0.8438 - val_accuracy: 0.7825
Epoch 2/5
6/6 [==============================] - 0s 34ms/step - loss: 0.7913 - accuracy: 0.7813 - val_loss: 0.3423 - val_accuracy: 0.9102
Epoch 3/5
6/6 [==============================] - 0s 34ms/step - loss: 0.3289 - accuracy: 0.9096 - val_loss: 0.3661 - val_accuracy: 0.8785
Epoch 4/5
6/6 [==============================] - 0s 33ms/step - loss: 0.2984 - accuracy: 0.9070 - val_loss: 0.5444 - val_accuracy: 0.8468
Epoch 5/5
6/6 [==============================] - 0s 33ms/step - loss: 0.5225 - accuracy: 0.8662 - val_loss: 0.1936 - val_accuracy: 0.9465
313/313 - 0s - loss: 0.1935 - accuracy: 0.9432
Test loss: 0.19351395964622498
Test accuracy: 0.9431999921798706

モデルはサンプルあたりの時間をログに出力します（例:「4us / sample」）。 TensorFlow はモデルの最適化にある程度の時間を費やすため、最初のエポックは遅くなる可能性がありますが、その後はサンプルあたりの時間が安定するはずです。

このガイドを Colab で実行している場合は、混合精度と float32 のパフォーマンスの比較ができます。これを行うには、「dtype ポリシーを設定する」のセクションに従ってポリシーをmixed_float16からfloat32に変更し、この時点までのすべてのセルを再実行します。コンピューティング機能が 7.0 以上の GPU では、サンプルあたりの時間が大幅に増加し、混合精度がモデルを高速化していることが分かるはずです。 例えば Titan V GPU の場合、サンプルあたりの時間が 4us から 12us に増加します。ガイドを続行する前に、必ずポリシーをmixed_float16に戻し、セルを再実行してください。

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

TPU で混合精度を実行する場合、GPU で混合精度を実行する場合に比べ、パフォーマンスはそれほど向上しません。これは、dtype ポリシーのデフォルトがfloat32であっても、TPU が内部で bfloat16 を使用して特定の操作を実行しているためです。TPUハードウェアは、matmul など bfloat16 で数値的に安定している特定の ops（オプス）の float32 には対応していません。そのような ops には、TPU バックエンドが代わりに内部的に静かに bfloat16 を使用します。結果として、そのような ops を使用するレイヤーにdtype='float32'を渡しても数値的な効果はありませんが、このようなレイヤーに bfloat16 計算を実行することが有害となることもありません。

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

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$）を掛けます。この数値を loss scale（損失スケール）と呼びます。これによって、勾配も $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=0, 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)

WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/mixed_precision/loss_scale.py:54: FixedLossScale.__init__ (from tensorflow.python.training.experimental.loss_scale) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.keras.mixed_precision.LossScaleOptimizer instead. LossScaleOptimizer now has all the functionality of FixedLossScale
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 のみで実行する場合に比べ、2 つの変更が必要です。

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

手順 (2) では、tf.keras.mixed_precision.experimental.LossScaleOptimizer クラスを使用し、オプティマイザをラップして損失スケーリングを適用します。これにはオプティマイザと損失スケールの 2 つの引数が必要です。動的損失スケールを使用するには、次のように構成します。

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

WARNING:tensorflow:tf.keras.mixed_precision.experimental.LossScaleOptimizer is deprecated. Please use tf.keras.mixed_precision.LossScaleOptimizer instead. For example
  opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt)

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

次に、トレーニングステップ関数を定義します。 損失をスケーリングし、勾配のスケーリングを解除するために、損失スケールオプティマイザの 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=4.1430134773254395, test accuracy=0.6510000228881836
Epoch 1: loss=0.5511777400970459, test accuracy=0.8744999766349792
Epoch 2: loss=0.32883334159851074, test accuracy=0.9394999742507935
Epoch 3: loss=0.306718111038208, test accuracy=0.9509000182151794
Epoch 4: loss=0.23808135092258453, test accuracy=0.8967999815940857

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

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

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

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

GPU の Tensor Core を使用する

前述したように、最近の NVIDIA GPU は、float16 行列を非常に速く乗算できる Tensor Core と呼ばれる特殊なハードウェアユニットを使用しています。ただし Tensor Core は、テンソルの特定の次元を 8 の倍数にする必要があります。以下の例では、Tensor Core を使用するために引数を 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)

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

XLA

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

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

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

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

まとめ

  • TPU または 7.0 以上のコンピューティング機能を備えた 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を呼び出して勾配をスケーリング解除する。

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

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

tf.keras.mixed_precisionAPI を使用した混合精度の例をもっとご覧になりたい方は、公式モデルレポジトリ をご参照ください。ResNet や Transformer など大部分の公式モデルは、--dtype=fp16を渡すと混合精度を使用して実行します。