#@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.

"""
Explanation: Copyright 2019 The TensorFlow Authors.
End of explanation
"""


# Import TensorFlow and TensorFlow Datasets

import tensorflow_datasets as tfds
import tensorflow as tf
tfds.disable_progress_bar()

import os

print(tf.__version__)

"""
Explanation: Keras による分散型トレーニング
<table class="tfo-notebook-buttons" align="left">
  <td>
<img src="https://www.tensorflow.org/images/tf_logo_32px.png"><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/tutorials/distribute/keras.ipynb">TensorFlow.org で表示</a>
</td>
  <td>
<img src="https://www.tensorflow.org/images/colab_logo_32px.png"><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ja/tutorials/distribute/keras.ipynb">Google Colab で実行</a>
</td>
  <td>
<img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png"><a target="_blank" href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/tutorials/distribute/keras.ipynb">GitHub でソースを表示</a>
</td>
  <td>
<img src="https://www.tensorflow.org/images/download_logo_32px.png"><a href="https://storage.googleapis.com/tensorflow_docs/docs/site/en/tutorials/distribute/keras.ipynb">ノートブックをダウンロード</a>
</td>
</table>

概要
tf.distribute.Strategy API は、複数の処理ユニットに渡ってトレーニングを分散するための抽象化を提供します。目標は、ユーザーが既存のモデルとトレーニングコードを使用して、最小限の変更で分散型トレーニングをできるようにすることです。
このチュートリアルでは、tf.distribute.MirroredStrategyを使用します。これは、1 台のマシン上の多数の GPU で同期トレーニングを行うグラフ内レプリケーションを行います。基本的には、モデルの変数の全ての各プロセッサにコピーします。その後、all-reduceを使用して全プロセッサからの勾配を結合することで、結合された値をモデルの全コピーに適用します。
MirroredStrategyは、TensorFlow コアで利用可能な数ある分散ストラテジーの中の 1 つです。その他のストラテジーについては、分散ストラテジーガイドをご覧ください。
Keras API
この例では、tf.keras API を使用してモデルとトレーニングループを構築します。カスタムトレーニングループについては、トレーニングループを使用した tf.distribute.Strategy チュートリアルをご覧ください。
依存関係をインポートする
End of explanation
"""


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

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

"""
Explanation: データセットをダウンロードする
MNIST データセットをダウンロードして TensorFlow Datasets から 読み込みます。これはtf.data形式のデータセットを返します。
with_infoをTrueに設定すると、データセット全体に対するメタデータを含み、ここではinfoに保存されます。このメタデータオブジェクトは、トレーニングとテストの例の数を含みます。
End of explanation
"""


strategy = tf.distribute.MirroredStrategy()

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

"""
Explanation: 分散ストラテジーを定義する
MirroredStrategyオブジェクトを作成します。これは分散を処理し、コンテキストマネージャ（tf.distribute.MirroredStrategy.scope）を提供して内側にモデルを構築します。
End of explanation
"""


# 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

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


def scale(image, label):
  image = tf.cast(image, tf.float32)
  image /= 255

  return image, label

"""
Explanation: 0-255 のピクセル値は 0-1 の範囲に正規化する必要があります。このスケールを関数で定義します。
End of explanation
"""


train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)

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


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

"""
Explanation: モデルを作成する
strategy.scopeのコンテキストで Keras モデルを作成し、コンパイルします。
End of explanation
"""


# 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 &lt; 3:
    return 1e-3
  elif epoch &gt;= 3 and epoch &lt; 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()
]

"""
Explanation: コールバックを定義する
ここでは以下のコールバックを使用します。

TensorBoard: このコールバックはグラフの可視化を可能にする TensorBoard のためのログを書き込みます。
モデルチェックポイント<br>: このコールバックは各エポック後のモデルを保存します。
学習率スケジューラ: このコールバックを使用すると、各エポック/バッチ後に変化する学習率をスケジュールすることができます。

例証目的として、このノートブックでは出力コールバックを追加して学習率を表示します。
End of explanation
"""


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

"""
Explanation: トレーニングして評価する
では、モデル上でfitを呼び出し、チュートリアルの最初に作成したデータセットを渡す、通常の方法でモデルをトレーニングします。トレーニングの分散の有無に関わらず、このステップは同じです。
End of explanation
"""


# check the checkpoint directory
!ls {checkpoint_dir}

"""
Explanation: 以下でご覧いただけるように、チェックポイントは保存されています。
End of explanation
"""


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

"""
Explanation: モデルのパフォーマンスを確認するには、最新のチェックポイントを読み込み、テストデータ上でevaluateを呼び出します。
適切なデータセットを使用して、前と同様にevaluateを呼び出します。
End of explanation
"""


!ls -sh ./logs

"""
Explanation: 出力の確認には、TensorBoard のログをダウンロードして端末に表示することができます。
$ tensorboard --logdir=path/to/log-directory
End of explanation
"""


path = 'saved_model/'

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

"""
Explanation: SavedModel にエクスポートする
グラフと変数をプラットフォームに依存しない SavedModel 形式にエクスポートします。モデルが保存された後、スコープの有無に関わらずそれを読み込むことができます。
End of explanation
"""


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

"""
Explanation: strategy.scopeなしでモデルを読み込みます。
End of explanation
"""


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

"""
Explanation: strategy.scopeと共にモデルを読み込みます。
End of explanation
"""


