![]() | ![]() | ![]() | ![]() |
概要
tf.distribute.Strategy
APIは、トレーニングを複数の処理ユニットに分散するための抽象化を提供します。これにより、最小限の変更で既存のモデルとトレーニングコードを使用して分散トレーニングを実行できます。
このチュートリアルでは、 tf.distribute.MirroredStrategy
を使用して、1台のマシン上の多数のGPUで同期トレーニングを使用してグラフ内レプリケーションを実行する方法を示します。この戦略は基本的に、モデルのすべての変数を各プロセッサーにコピーします。次に、 all-reduceを使用してすべてのプロセッサからの勾配を結合し、結合された値をモデルのすべてのコピーに適用します。
tf.keras
APIを使用してモデルを構築し、 Model.fit
を使用してトレーニングします。 (カスタムトレーニングループとMirroredStrategy
を使用した分散トレーニングについては、このチュートリアルをご覧ください。)
MirroredStrategy
は、単一のマシン上の複数のGPUでモデルをトレーニングします。複数のワーカーの多くのGPUで同期トレーニングを行うには、KerasModel.fitまたはカスタムトレーニングループでtf.distribute.MultiWorkerMirroredStrategy
を使用します。その他のオプションについては、分散トレーニングガイドを参照してください。
他のさまざまな戦略について学ぶために、TensorFlowを使用した分散トレーニングガイドがあります。
設定
import tensorflow_datasets as tfds
import tensorflow as tf
import os
# Load the TensorBoard notebook extension.
%load_ext tensorboard
print(tf.__version__)
2.8.0-rc1
データセットをダウンロードする
TensorFlowデータセットからMNISTデータセットをロードします。これにより、 tf.data
形式のデータセットが返されます。
with_info
引数をTrue
に設定すると、データセット全体のメタデータが含まれ、ここでinfo
に保存されます。特に、このメタデータオブジェクトには、トレインとテストの例の数が含まれています。
datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)
mnist_train, mnist_test = datasets['train'], datasets['test']
流通戦略を定義する
MirroredStrategy
オブジェクトを作成します。これは配布を処理し、モデルを内部に構築するためのコンテキストマネージャー( MirroredStrategy.scope
)を提供します。
strategy = tf.distribute.MirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',) INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
Number of devices: 1
入力パイプラインを設定します
複数のGPUを使用してモデルをトレーニングする場合、バッチサイズを大きくすることで、追加の計算能力を効果的に使用できます。一般に、GPUメモリに適合する最大のバッチサイズを使用し、それに応じて学習率を調整します。
# You can also do info.splits.total_num_examples to get the total
# number of examples in the dataset.
num_train_examples = info.splits['train'].num_examples
num_test_examples = info.splits['test'].num_examples
BUFFER_SIZE = 10000
BATCH_SIZE_PER_REPLICA = 64
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
画像のピクセル値を[0, 255]
の範囲から[0, 1]
]の範囲に正規化する関数を定義します(機能スケーリング)。
def scale(image, label):
image = tf.cast(image, tf.float32)
image /= 255
return image, label
このscale
関数をトレーニングデータとテストデータに適用してから、 tf.data.Dataset
APIを使用してトレーニングデータをシャッフルし( Dataset.shuffle
)、バッチ処理します( Dataset.batch
)。パフォーマンスを向上させるために、トレーニングデータのメモリ内キャッシュも保持していることに注意してください( Dataset.cache
)。
train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)
モデルを作成する
Strategy.scope
のコンテキストでKerasモデルを作成してコンパイルします。
with strategy.scope():
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10)
])
model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.Adam(),
metrics=['accuracy'])
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
コールバックを定義する
次のtf.keras.callbacks
を定義します。
-
tf.keras.callbacks.TensorBoard
:TensorBoardのログを書き込みます。これにより、グラフを視覚化できます。 -
tf.keras.callbacks.ModelCheckpoint
:すべてのエポックの後など、特定の頻度でモデルを保存します。 -
tf.keras.callbacks.LearningRateScheduler
:たとえば、すべてのエポック/バッチの後に変更する学習率をスケジュールします。
説明のために、 PrintLR
と呼ばれるカスタムコールバックを追加して、ノートブックに学習率を表示します。
# Define the checkpoint directory to store the checkpoints.
checkpoint_dir = './training_checkpoints'
# Define the name of the checkpoint files.
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")
# Define a function for decaying the learning rate.
# You can define any decay function you need.
def decay(epoch):
if epoch < 3:
return 1e-3
elif epoch >= 3 and epoch < 7:
return 1e-4
else:
return 1e-5
# Define a callback for printing the learning rate at the end of each epoch.
class PrintLR(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
print('\nLearning rate for epoch {} is {}'.format(epoch + 1,
model.optimizer.lr.numpy()))
# Put all the callbacks together.
callbacks = [
tf.keras.callbacks.TensorBoard(log_dir='./logs'),
tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix,
save_weights_only=True),
tf.keras.callbacks.LearningRateScheduler(decay),
PrintLR()
]
トレーニングと評価
次に、モデルでModel.fit
を呼び出し、チュートリアルの最初に作成したデータセットを渡すことにより、通常の方法でモデルをトレーニングします。この手順は、トレーニングを配布するかどうかに関係なく同じです。
EPOCHS = 12
model.fit(train_dataset, epochs=EPOCHS, callbacks=callbacks)
2022-01-26 05:38:28.865380: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:547] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. Epoch 1/12 INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',). 933/938 [============================>.] - ETA: 0s - loss: 0.2029 - accuracy: 0.9399 Learning rate for epoch 1 is 0.0010000000474974513 938/938 [==============================] - 10s 4ms/step - loss: 0.2022 - accuracy: 0.9401 - lr: 0.0010 Epoch 2/12 930/938 [============================>.] - ETA: 0s - loss: 0.0654 - accuracy: 0.9813 Learning rate for epoch 2 is 0.0010000000474974513 938/938 [==============================] - 3s 3ms/step - loss: 0.0652 - accuracy: 0.9813 - lr: 0.0010 Epoch 3/12 931/938 [============================>.] - ETA: 0s - loss: 0.0453 - accuracy: 0.9864 Learning rate for epoch 3 is 0.0010000000474974513 938/938 [==============================] - 3s 3ms/step - loss: 0.0453 - accuracy: 0.9864 - lr: 0.0010 Epoch 4/12 923/938 [============================>.] - ETA: 0s - loss: 0.0246 - accuracy: 0.9933 Learning rate for epoch 4 is 9.999999747378752e-05 938/938 [==============================] - 3s 3ms/step - loss: 0.0244 - accuracy: 0.9934 - lr: 1.0000e-04 Epoch 5/12 929/938 [============================>.] - ETA: 0s - loss: 0.0211 - accuracy: 0.9944 Learning rate for epoch 5 is 9.999999747378752e-05 938/938 [==============================] - 3s 3ms/step - loss: 0.0212 - accuracy: 0.9944 - lr: 1.0000e-04 Epoch 6/12 930/938 [============================>.] - ETA: 0s - loss: 0.0192 - accuracy: 0.9950 Learning rate for epoch 6 is 9.999999747378752e-05 938/938 [==============================] - 3s 3ms/step - loss: 0.0194 - accuracy: 0.9950 - lr: 1.0000e-04 Epoch 7/12 927/938 [============================>.] - ETA: 0s - loss: 0.0179 - accuracy: 0.9953 Learning rate for epoch 7 is 9.999999747378752e-05 938/938 [==============================] - 3s 3ms/step - loss: 0.0179 - accuracy: 0.9953 - lr: 1.0000e-04 Epoch 8/12 938/938 [==============================] - ETA: 0s - loss: 0.0153 - accuracy: 0.9966 Learning rate for epoch 8 is 9.999999747378752e-06 938/938 [==============================] - 3s 3ms/step - loss: 0.0153 - accuracy: 0.9966 - lr: 1.0000e-05 Epoch 9/12 927/938 [============================>.] - ETA: 0s - loss: 0.0151 - accuracy: 0.9966 Learning rate for epoch 9 is 9.999999747378752e-06 938/938 [==============================] - 3s 3ms/step - loss: 0.0150 - accuracy: 0.9966 - lr: 1.0000e-05 Epoch 10/12 935/938 [============================>.] - ETA: 0s - loss: 0.0148 - accuracy: 0.9966 Learning rate for epoch 10 is 9.999999747378752e-06 938/938 [==============================] - 3s 3ms/step - loss: 0.0148 - accuracy: 0.9966 - lr: 1.0000e-05 Epoch 11/12 937/938 [============================>.] - ETA: 0s - loss: 0.0146 - accuracy: 0.9967 Learning rate for epoch 11 is 9.999999747378752e-06 938/938 [==============================] - 3s 3ms/step - loss: 0.0146 - accuracy: 0.9967 - lr: 1.0000e-05 Epoch 12/12 926/938 [============================>.] - ETA: 0s - loss: 0.0145 - accuracy: 0.9967 Learning rate for epoch 12 is 9.999999747378752e-06 938/938 [==============================] - 3s 3ms/step - loss: 0.0144 - accuracy: 0.9967 - lr: 1.0000e-05 <keras.callbacks.History at 0x7fad70067c10>プレースホルダー19
保存されたチェックポイントを確認します。
# Check the checkpoint directory.
ls {checkpoint_dir}
checkpoint ckpt_4.data-00000-of-00001 ckpt_1.data-00000-of-00001 ckpt_4.index ckpt_1.index ckpt_5.data-00000-of-00001 ckpt_10.data-00000-of-00001 ckpt_5.index ckpt_10.index ckpt_6.data-00000-of-00001 ckpt_11.data-00000-of-00001 ckpt_6.index ckpt_11.index ckpt_7.data-00000-of-00001 ckpt_12.data-00000-of-00001 ckpt_7.index ckpt_12.index ckpt_8.data-00000-of-00001 ckpt_2.data-00000-of-00001 ckpt_8.index ckpt_2.index ckpt_9.data-00000-of-00001 ckpt_3.data-00000-of-00001 ckpt_9.index ckpt_3.index
モデルのパフォーマンスを確認するには、最新のチェックポイントをロードし、テストデータModel.evaluate
を呼び出します。
model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
eval_loss, eval_acc = model.evaluate(eval_dataset)
print('Eval loss: {}, Eval accuracy: {}'.format(eval_loss, eval_acc))
2022-01-26 05:39:15.260539: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:547] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. 157/157 [==============================] - 2s 4ms/step - loss: 0.0373 - accuracy: 0.9879 Eval loss: 0.03732967749238014, Eval accuracy: 0.9879000186920166
出力を視覚化するには、TensorBoardを起動し、ログを表示します。
%tensorboard --logdir=logs
ls -sh ./logs
total 4.0K 4.0K trainプレースホルダー26
SavedModelにエクスポート
Model.saveを使用して、グラフと変数をプラットフォームに依存しないModel.save
形式にエクスポートします。モデルを保存した後、 Strategy.scope
の有無にかかわらずモデルをロードできます。
path = 'saved_model/'
model.save(path, save_format='tf')
2022-01-26 05:39:18.012847: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them. INFO:tensorflow:Assets written to: saved_model/assets INFO:tensorflow:Assets written to: saved_model/assets
ここで、 Strategy.scope
なしでモデルをロードします。
unreplicated_model = tf.keras.models.load_model(path)
unreplicated_model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.Adam(),
metrics=['accuracy'])
eval_loss, eval_acc = unreplicated_model.evaluate(eval_dataset)
print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
157/157 [==============================] - 1s 2ms/step - loss: 0.0373 - accuracy: 0.9879 Eval loss: 0.03732967749238014, Eval Accuracy: 0.9879000186920166
Strategy.scope
を使用してモデルをロードします。
with strategy.scope():
replicated_model = tf.keras.models.load_model(path)
replicated_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.Adam(),
metrics=['accuracy'])
eval_loss, eval_acc = replicated_model.evaluate(eval_dataset)
print ('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))
2022-01-26 05:39:19.489971: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:547] The `assert_cardinality` transformation is currently not handled by the auto-shard rewrite and will be removed. 157/157 [==============================] - 3s 3ms/step - loss: 0.0373 - accuracy: 0.9879 Eval loss: 0.03732967749238014, Eval Accuracy: 0.9879000186920166
追加のリソース
Keras Model.fit
APIでさまざまな配布戦略を使用するその他の例:
- TPUチュートリアルでBERTを使用してGLUEタスクを解決するには、GPUでのトレーニングに
tf.distribute.MirroredStrategy
を使用し、TPUでのtf.distribute.TPUStrategy
を使用します。 - 配布戦略チュートリアルを使用したモデルの保存と読み込みでは、tf.distribute.Strategyで
tf.distribute.Strategy
を使用する方法を示します。 - 公式のTensorFlowモデルは、複数の配信戦略を実行するように構成できます。
TensorFlow配信戦略の詳細については、以下をご覧ください。
- tf.distribute.Strategyチュートリアルを使用したカスタムトレーニングは、カスタムトレーニングループを使用したシングルワーカートレーニングに
tf.distribute.MirroredStrategy
を使用する方法を示しています。 - Kerasチュートリアルを使用したMulti-workerトレーニングは、
Model.fit
でMultiWorkerMirroredStrategy
を使用する方法を示しています。 - KerasとMultiWorkerMirroredStrategyチュートリアルを使用したカスタムトレーニングループは、Kerasとカスタムトレーニングループを使用して
MultiWorkerMirroredStrategy
を使用する方法を示しています。 - TensorFlowガイドの分散トレーニングでは、利用可能な分散戦略の概要を説明します。
- tf.functionガイドによるパフォーマンスの向上は、TensorFlowモデルのパフォーマンスを最適化するために使用できるTensorFlow Profilerなど、他の戦略やツールに関する情報を提供します。