GitHub Repository: tensorflow/docs-l10n
Path: blob/master/site/ja/tutorials/distribute/keras.ipynb
²⁵¹¹⁸ views

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

In [ ]:

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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

概要

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

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

tf.keras API を使用して、モデルとそれをトレーニングするための Model.fit をビルドします。（カスタムトレーニングループと MirroredStrategy を使った分散型トレーニングについては、こちらのチュートリアルをご覧ください。）

MirroredStrategy は単一のマシンの複数の GPU でモデルをトレーニングします。複数のワーカーの多数の GPU で同期トレーニングを行う場合は、tf.distribute.MultiWorkerMirroredStrategy とKeras の Model.fit かカスタムトレーニングループを使用します。その他のオプションについては、分散型トレーニングガイドをご覧ください。

その他のさまざまなストラテジーについては、TensorFlow の分散型トレーニングガイドをご覧ください。

セットアップ

In [ ]:

import tensorflow_datasets as tfds
import tensorflow as tf

import os

# Load the TensorBoard notebook extension.
%load_ext tensorboard

In [ ]:

print(tf.__version__)

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

TensorFlow Datasets から MNIST データセットを読み込みます。これは、tf.data 形式のデータセットを返します。

with_info 引数を True に設定すると、データセット全体に対するメタデータが含まれます。ここでは info に保存されます。このメタデータオブジェクトには、トレーニングとテストの例の数などが含まれます。

In [ ]:

datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)

mnist_train, mnist_test = datasets['train'], datasets['test']

分散ストラテジーを定義する

MirroredStrategy オブジェクトを作成します。これは分散を処理し、モデル内に構築するコンテキストマネージャ (MirroredStrategy.scope) を提供します。

In [ ]:

strategy = tf.distribute.MirroredStrategy()

In [ ]:

print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

入力パイプラインをセットアップする

マルチ GPU でモデルをトレーニングする場合、バッチサイズを増加することにより追加の計算能力を効果的に利用することができます。一般的には、GPU メモリに収まる最大のバッチサイズを使用し、それに応じて学習率を調整します。

In [ ]:

# 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] の範囲に正規化する関数を定義します（特徴量スケーリング）。

In [ ]:

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）。

In [ ]:

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 API を使ってモデルを作成し、コンパイルします。

In [ ]:

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(learning_rate=0.001),
                metrics=['accuracy'])

この MNIST データセットを使ったトイサンプルでは、Adam オプティマイザのデフォルトの学習率である 0.001 を使用します。

より大規模なデータセットの場合、分散トレーニングの主なメリットはトレーニングステップごとにより多くの学習を行えることです。これは、各ステップがより多くのトレーニングデータを並行して処理するため、（モデルとデータセットの制限内で）より大きな学習率が可能となるためです。

コールバックを定義する

以下の Keras コールバックを定義します。

tf.keras.callbacks.TensorBoard: グラフを視覚化できるように、TensorBoard 用のログを書き込みます。
tf.keras.callbacks.ModelCheckpoint: 各エポック後など、特定の頻度でモデルを保存します。
tf.keras.callbacks.BackupAndRestore: モデルと現在のエポック番号をバックアップすることで、フォールトトレランス機能を提供します。詳細は、Keras によるマルチワーカートレーニングチュートリアルのフォールトトレランスセクションをご覧ください。
tf.keras.callbacks.LearningRateScheduler: schedules the learning rate to change after, for example, every epoch/batch.

このノートブックでは例示目的で、PrintLR というカスタムコールバックを追加して、学習率を表示します。

注意: ジョブの失敗から再開する際に、トレーニング状態をリストアするための主なメカニズムとして、ModelCheckpoint の代わりに BackupAndRestore コールバックを使用してください。BackupAndRestore は eager モードのみをサポートするため、graph モードでは ModelCheckpoint を使用することを検討してください。

In [ ]:

# 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}")

In [ ]:

# 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

In [ ]:

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

In [ ]:

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

トレーニングして評価する

次に、通常の方法でモデルをトレーニングします。モデル上で Keras Model.fit を呼び出し、チュートリアルの最初に作成したデータセットを渡します。トレーニングを分散しているかに関わらず、このステップは同じです。

In [ ]:

EPOCHS = 12

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

保存済みのチェックポイントを確認します。

In [ ]:

# Check the checkpoint directory.
!ls {checkpoint_dir}

モデルがどれほどうまく実行するかを確認するために、最新のチェックポイントを読み込み、テストデータで Model.evaluate を呼び出します。

In [ ]:

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

出力を視覚化するために、TensorBoard を起動して、ログを表示します。

In [ ]:

%tensorboard --logdir=logs

In [ ]:

!ls -sh ./logs

モデルを保存する

Model.save を使用して、モデルを .keras zip アーカイブに保存します。モデルが保存されたら、Strategy.scope の有無に関係なくそれを読み込めるようになります。

In [ ]:

path = 'my_model.keras'

In [ ]:

model.save(path)

次に、Strategy.scope を使用せずにモデルを読み込みます。

In [ ]:

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

Strategy.scope を使用してモデルを読み込みます。

In [ ]:

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

追加リソース

さまざまな分散ストラテジーと Keras Model.fit API を使用したその他の例をご覧ください。

TPU で BERT を使って GLUE タスクを解決するチュートリアルでは、GPU でのトレーニングには tf.distribute.MirroredStrategy を使用し、TPU では tf.distribute.TPUStrategy を使用しています。
分散ストラテジーを使ってモデルを保存して読み込むチュートリアルでは、SavedModel API と tf.distribute.Strategy の使用方法が説明されています。
TensorFlow 公式モデルは、複数の分散ストラテジーを実行できるように構成可能です。

TensorFlow 分散ストラテジーに関してさらに学習するには、以下をご覧ください。

tf.distribute.Strategy によるカスタムトレーニングチュートリアルでは、カスタムトレーニングループを使って単一ワーカートレーニングに tf.distribute.MirroredStrategy を使用する方法が説明されています。
Keras によるマルチワーカートレーニングのチュートリアルでは、MultiWorkerMirroredStrategy と Model.fit を使用する方法が説明されています。
Keras によるカスタムトレーニングループと MultiWorkerMirroredStrategy のチュートリアルでは、Keras とカスタムトレーニングループでMultiWorkerMirroredStrategy を使用する方法が説明されています。
TensorFlow での分散型トレーニングガイドでは、利用可能な分散ストラテジーの概要が説明されています。
tf.function を使ったパフォーマンスの改善ガイドでは、その他のストラテジーや、TensorFlow モデルのパフォーマンスを最適化するために使用できる TensorFlow Profiler といったツールに関する情報が提供されています。

注意: tf.distribute.Strategy の開発は積極積に進められています。近日中にはより多くの例やチュートリアルを追加する予定ですので、ぜひお試しください。フィードバックをお待ちしております。GitHub の課題から、お気軽にお寄せください。