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

混合精度

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

概観

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

現在、ほとんどのモデルで32ビットのメモリを使用するfloat32 dtypeを使用しています。ただし、精度の低い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-a0cb13c6-7f41-da14-7d18-cd7349a7ddc8)

すべての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)
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

ポリシーは、レイヤーの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

この例では、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.1244 - accuracy: 0.3937 - val_loss: 0.8114 - val_accuracy: 0.8036
Epoch 2/5
6/6 [==============================] - 0s 26ms/step - loss: 0.6696 - accuracy: 0.8071 - val_loss: 0.3220 - val_accuracy: 0.9084
Epoch 3/5
6/6 [==============================] - 0s 26ms/step - loss: 0.3986 - accuracy: 0.8744 - val_loss: 0.3270 - val_accuracy: 0.8909
Epoch 4/5
6/6 [==============================] - 0s 25ms/step - loss: 0.2740 - accuracy: 0.9110 - val_loss: 0.1909 - val_accuracy: 0.9448
Epoch 5/5
6/6 [==============================] - 0s 25ms/step - loss: 0.3640 - accuracy: 0.8932 - val_loss: 0.4617 - val_accuracy: 0.8450
313/313 - 1s - loss: 0.4765 - accuracy: 0.8403
Test loss: 0.47645366191864014
Test accuracy: 0.8403000235557556

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

このガイドをColabで実行している場合は、混合精度とfloat32のパフォーマンスを比較できます。そのためには、「dtypeポリシーの設定」セクションでポリシーをfloat32からmixed_float16float32し、この時点までのすべてのセルを再実行します。少なくともコンピューティング機能が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の計算でそのようなレイヤーを実行することは有害です。

損失のスケーリング

損失スケーリングは技術であるtf.keras.Model.fitで自動的に行いをmixed_float16数値アンダーフローを避けるために、ポリシー。このセクションでは、損失のスケーリングとその動作をカスタマイズする方法について説明します。

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

float16データ型は、float32と比較してダイナミックレンジが狭い。つまり、$ 65504 $を超える値はオーバーフローして無限大になり、$ 6.0を下回る値は10倍にアンダーフローしてゼロになります。 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オブジェクトが関連付けられて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=3.38541841506958, test accuracy=0.6798999905586243
Epoch 1: loss=0.5135143399238586, test accuracy=0.8654999732971191
Epoch 2: loss=0.3303307294845581, test accuracy=0.9258000254631042
Epoch 3: loss=0.19915412366390228, test accuracy=0.9215999841690063
Epoch 4: loss=0.2938493490219116, test accuracy=0.9502000212669373

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を使用した混合精度のその他の例については、 公式のモデルリポジトリをご覧ください。 ResNetTransformerなどのほとんどの公式モデルは、-- --dtype=fp16を渡すことにより、混合精度を使用して実行されます。