Google I / Oの基調講演、製品セッション、ワークショップなどを見るプレイリストを見る

Kerasによる分散トレーニング

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

概要概要

tf.distribute.Strategy APIは、トレーニングを複数の処理ユニットに分散するための抽象化を提供します。目標は、ユーザーが最小限の変更で既存のモデルとトレーニングコードを使用して分散トレーニングを有効にできるようにすることです。

このチュートリアルでは、 tf.distribute.MirroredStrategy使用しtf.distribute.MirroredStrategy 。これは、1台のマシン上の多数のGPUで同期トレーニングを使用してグラフ内レプリケーションを実行します。基本的に、モデルのすべての変数を各プロセッサにコピーします。次に、 all-reduceを使用してすべてのプロセッサからの勾配を結合し、結合された値をモデルのすべてのコピーに適用します。

MirroredStrategyは、TensorFlowコアで利用できるいくつかの配信戦略の1つです。あなたは流通戦略ガイドでより多くの戦略について読むことができます。

Keras API

この例では、 tf.keras APIを使用して、モデルとトレーニングループを構築します。カスタムトレーニングループについては、トレーニングループを使用したtf.distribute.Strategyチュートリアルを参照してください。

依存関係をインポートする

# Import TensorFlow and TensorFlow Datasets

import tensorflow_datasets as tfds
import tensorflow as tf

import os

print(tf.__version__)

2.5.0

データセットをダウンロードする

MNISTデータセットをダウンロードしてからそれをロードするTensorFlowデータセット。これにより、データセットがtf.data形式で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オブジェクトを作成します。これは配布を処理し、内部にモデルを構築するためのコンテキストマネージャー（ tf.distribute.MirroredStrategy.scope ）を提供します。

strategy = tf.distribute.MirroredStrategy()

WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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

この関数をトレーニングおよびテストデータに適用し、トレーニングデータをシャッフルして、トレーニング用にバッチ処理します。パフォーマンスを向上させるために、トレーニングデータのメモリ内キャッシュも保持していることに注意してください。

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

コールバックを定義する

ここで使用されるコールバックは次のとおりです。

  • TensorBoard ：このコールバックは、グラフを視覚化できるTensorBoardのログを書き込みます。
  • モデルチェックポイント：このコールバックは、エポックごとにモデルを保存します。
  • 学習率スケジューラ：このコールバックを使用して、エポック/バッチごとに変更する学習率をスケジュールできます。

説明のために、印刷コールバックを追加して、ノートブックに学習率を表示します。

# Define the checkpoint directory to store the checkpoints

checkpoint_dir = './training_checkpoints'
# Name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

# 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

# Callback for printing the LR 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()))

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()
]

トレーニングと評価

次に、通常の方法でモデルをトレーニングし、モデルでfitを呼び出し、チュートリアルの最初に作成したデータセットを渡します。この手順は、トレーニングを配布するかどうかに関係なく同じです。

model.fit(train_dataset, epochs=12, callbacks=callbacks)

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',).
3/938 [..............................] - ETA: 4:06 - loss: 2.2662 - accuracy: 0.2292WARNING:tensorflow:Callback method `on_train_batch_begin` is slow compared to the batch time (batch time: 0.0044s vs `on_train_batch_begin` time: 0.0716s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_begin` is slow compared to the batch time (batch time: 0.0044s vs `on_train_batch_begin` time: 0.0716s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0044s vs `on_train_batch_end` time: 0.0148s). Check your callbacks.
WARNING:tensorflow:Callback method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0044s vs `on_train_batch_end` time: 0.0148s). Check your callbacks.
938/938 [==============================] - 9s 4ms/step - loss: 0.2016 - accuracy: 0.9415

Learning rate for epoch 1 is 0.0010000000474974513
Epoch 2/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0685 - accuracy: 0.9791

Learning rate for epoch 2 is 0.0010000000474974513
Epoch 3/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0477 - accuracy: 0.9857

Learning rate for epoch 3 is 0.0010000000474974513
Epoch 4/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0273 - accuracy: 0.9925

Learning rate for epoch 4 is 9.999999747378752e-05
Epoch 5/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0236 - accuracy: 0.9937

Learning rate for epoch 5 is 9.999999747378752e-05
Epoch 6/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0219 - accuracy: 0.9942

Learning rate for epoch 6 is 9.999999747378752e-05
Epoch 7/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0201 - accuracy: 0.9949

Learning rate for epoch 7 is 9.999999747378752e-05
Epoch 8/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0176 - accuracy: 0.9958

Learning rate for epoch 8 is 9.999999747378752e-06
Epoch 9/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0173 - accuracy: 0.9959

Learning rate for epoch 9 is 9.999999747378752e-06
Epoch 10/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0171 - accuracy: 0.9959

Learning rate for epoch 10 is 9.999999747378752e-06
Epoch 11/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0169 - accuracy: 0.9959

Learning rate for epoch 11 is 9.999999747378752e-06
Epoch 12/12
938/938 [==============================] - 3s 3ms/step - loss: 0.0168 - accuracy: 0.9962

Learning rate for epoch 12 is 9.999999747378752e-06
<tensorflow.python.keras.callbacks.History at 0x7f74800d0550>

以下に示すように、チェックポイントは保存されています。

# 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

モデルのパフォーマンスを確認するには、最新のチェックポイントをロードし、テストデータでevaluateを呼び出します。

適切なデータセットを使用して、以前と同じように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))

157/157 [==============================] - 2s 4ms/step - loss: 0.0396 - accuracy: 0.9867
Eval loss: 0.03959868103265762, Eval Accuracy: 0.9866999983787537

出力を確認するには、ターミナルでTensorBoardログをダウンロードして表示できます。

$ tensorboard --logdir=path/to/log-directory

ls -sh ./logs

total 4.0K
4.0K train

SavedModelにエクスポート

グラフと変数をプラットフォームに依存しないSavedModel形式にエクスポートします。モデルを保存した後、スコープの有無にかかわらずモデルをロードできます。

path = 'saved_model/'

model.save(path, save_format='tf')

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 [==============================] - 0s 2ms/step - loss: 0.0396 - accuracy: 0.9867
Eval loss: 0.03959868103265762, Eval Accuracy: 0.9866999983787537

strategy.scope .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))

157/157 [==============================] - 3s 3ms/step - loss: 0.0396 - accuracy: 0.9867
Eval loss: 0.03959868103265762, Eval Accuracy: 0.9866999983787537

例とチュートリアル

keras fit / compileで配布戦略を使用する例を次に示します。

  1. tf.distribute.MirroredStrategyを使用してトレーニングされたTransformerの例
  2. tf.distribute.MirroredStrategyを使用してトレーニングされたNCFの例。

流通戦略ガイドに記載されているその他の例

次のステップ