このページは Cloud Translation API によって翻訳されました。
Switch to English

混合精度

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

概観

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

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

NVIDIA GPUはfloat32よりもfloat16で演算を実行でき、TPUはfloat32よりもbfloat16で演算を実行できます。したがって、これらのデバイスでは、可能な限り、これらの低精度のdtypeを使用する必要があります。ただし、変数といくつかの計算は、モデルが同じ品質にトレーニングされるように、数値上の理由からfloat32のままにする必要があります。 Keras混合精度APIを使用すると、float16またはbfloat16とfloat32の組み合わせを使用して、float16 / bfloat16によるパフォーマンスの利点とfloat32による数値の安定性の利点を得ることができます。

セットアップ

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コアと呼ばれる特別なハードウェアユニットを備えているため、混合精度のパフォーマンスが最大になります。古いGPUは、混合精度を使用するための数学パフォーマンスの利点を提供しませんが、メモリと帯域幅の節約により、いくつかの高速化が可能になります。 NVIDIAのCUDA GPU Webページで GPUのコンピューティング機能を調べることができます 。混合精度から最も恩恵を受けるGPUの例には、RTX GPU、Titan V、V100などがあります。

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

nvidia-smi -L

GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-843a0a0c-a559-ff09-842e-2a4fdb142480)

すべてのCloud TPUはbfloat16をサポートしています。

CPUや古いGPUでもスピードアップは期待されていませんが、単体テスト、デバッグ、または単にAPIを試すために、混合精度APIを使用できます。

dtypeポリシーの設定

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

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

ポリシーは、レイヤーの2つの重要な側面を指定します。レイヤーの計算が行われる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)
 
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

モデルの最後でのソフトマックスのアクティブ化は、float32にする必要があります。 dtypeポリシーはmixed_float16であるため、softmaxアクティベーションは通常、float16計算dtypeを持ち、float16テンソルを出力します。

これは、高密度レイヤーとdtype='float32'レイヤーを分離し、 dtype='float32'を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=mixed_precision.Policy('float32')渡すこともできます。レイヤーは常にdtype引数をポリシーに変換します。 Activationレイヤーには変数がないため、ポリシーの変数dtypeは無視されますが、ポリシーの計算dtype float32はsoftmaxを引き起こし、モデル出力はfloat32になります。

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

float16計算で数値的に安定しないと思われる場合は、 dtype='float32'を渡すことにより、任意のレイヤーのdtypeをfloat32にオーバーライドできます。ただし、ほとんどのレイヤーはmixed_float16およびmixed_bfloat16十分な精度を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
 
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step

この例では、int8からfloat32に入力データをキャストします。 255による除算はCPUで行われるため、float16にキャストしません。CPUは、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 [==============================] - 0s 55ms/step - loss: 4.4853 - accuracy: 0.4174 - val_loss: 0.8197 - val_accuracy: 0.7680
Epoch 2/5
6/6 [==============================] - 0s 28ms/step - loss: 0.7953 - accuracy: 0.7641 - val_loss: 0.4589 - val_accuracy: 0.8625
Epoch 3/5
6/6 [==============================] - 0s 27ms/step - loss: 0.3360 - accuracy: 0.9025 - val_loss: 0.3131 - val_accuracy: 0.8962
Epoch 4/5
6/6 [==============================] - 0s 27ms/step - loss: 0.3322 - accuracy: 0.8883 - val_loss: 0.2289 - val_accuracy: 0.9327
Epoch 5/5
6/6 [==============================] - 0s 27ms/step - loss: 0.2313 - accuracy: 0.9289 - val_loss: 0.2115 - val_accuracy: 0.9351
313/313 - 1s - loss: 0.2157 - accuracy: 0.9295
Test loss: 0.21573331952095032
Test accuracy: 0.9294999837875366

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

このガイドをColabで実行している場合は、混合精度とfloat32のパフォーマンスを比較できます。これを行うには、「dtypeポリシーの設定」セクションでポリシーをfloat32からmixed_float16に変更してから、この時点までのすべてのセルを再実行します。少なくともコンピューティング機能が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をサポートしていません。そのような操作の場合、TPUバックエンドは代わりに内部的にbfloat16を使用します。その結果、そのようなopsを使用するレイヤーにdtype='float32'を渡しても数値的な影響はないかもしれませんが、bfloat16の計算でそのようなレイヤーを実行することは有害です。

損失のスケーリング

損失スケーリングは技術である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
# 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=30, 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)
 
FixedLossScale(1024.0)

dtypeポリシーコンストラクターは常に損失スケールをLossScaleオブジェクトに変換します。この場合、それはtf.mixed_precision.experimental.FixedLossScaleに変換されます。これは、 DynamicLossScale以外の唯一のLossScaleサブクラスDynamicLossScale 。

レイヤーと同様に、モデルにはそれぞれdtypeポリシーがあります。存在する場合、モデルはポリシーの損失スケールを使用して、 tf.keras.Model.fitメソッドで損失スケーリングを適用します。つまり、 Model.fitが使用されている場合、損失のスケーリングについてまったく心配する必要はありません。デフォルトでは、 mixed_float16ポリシーには動的な損失スケールがあり、 Model.fitがそれを適用します。

カスタムトレーニングループでは、モデルはポリシーの損失スケールを無視するため、手動で適用する必要があります。これについては、次のセクションで説明します。

カスタムトレーニングループでモデルをトレーニングする

ここまでは、 tf.keras.Model.fitを使用して、混合精度でtf.keras.Model.fitモデルをトレーニングしtf.keras.Model.fit 。次に、カスタムトレーニングループで混合精度を使用します。カスタムトレーニングループがまだわからない場合は、まずカスタムトレーニングガイドをお読みください。

混合精度でカスタムトレーニングループを実行するには、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')

'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) ：スケーリングされた勾配のリストを入力として受け取り、それぞれを損失スケールで除算してスケールを解除します

勾配のアンダーフローを防ぐために、これらの関数を使用する必要があります。 LossScaleOptimizer.apply_gradientsは、InfsまたはNaNを持たないグラデーションを適用します。損失スケールも更新され、勾配にInfsまたは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.186450004577637, test accuracy=0.8429999947547913
Epoch 1: loss=0.5681900382041931, test accuracy=0.9300000071525574
Epoch 2: loss=0.330956369638443, test accuracy=0.9405999779701233
Epoch 3: loss=0.2619778513908386, test accuracy=0.9279000163078308
Epoch 4: loss=0.17072156071662903, test accuracy=0.9229999780654907

GPUパフォーマンスのヒント

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

バッチサイズを増やす

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

GPU Tensorコアが使用されていることを確認する

前述のように、最新のNVIDIA GPUは、float16行列を非常に迅速に乗算できるTensor Coresと呼ばれる特別なハードウェアユニットを使用しています。ただし、Tensorコアでは、テンソルの特定の次元を8の倍数にする必要があります。以下の例では、Tensorコアを使用するために引数が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コアを使用するようにしてください。詳細については、 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またはNVIDIA GPUを使用する場合は、少なくとも計算能力7.0で混合精度を使用する必要があります。これにより、パフォーマンスが最大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_precision APIを使用した混合精度の例については、 公式のモデルリポジトリをご覧ください。 ResNetやTransformerなどのほとんどの公式モデルは、-- --dtype=fp16を渡すことにより、混合精度を使用して実行されます。