#@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 2018 The TensorFlow Authors.
End of explanation
"""


import tensorflow as tf

import os
import tensorflow_datasets as tfds

"""
Explanation: TPU の使用
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://www.tensorflow.org/guide/tpu"><img src="https://www.tensorflow.org/images/tf_logo_32px.png">TensorFlow.org で表示</a></td>
  <td>     <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/guide/tpu.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png">Google Colab で実行</a>   </td>
  <td>     <a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ja/guide/tpu.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png">GitHub でソースを表示</a>   </td>
  <td>     <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/guide/tpu.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a>   </td>
</table>

この Colab ノートブックをダウンロードする前に、Runtime &gt; Change runtime type &gt; Hardware accelerator &gt; TPU でノートブックの設定を確認し、ハードウェアアクセラレータが TPU であることを確認してください。
セットアップ
End of explanation
"""


resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

"""
Explanation: TPU の初期化
TPU は通常 Cloud TPU ワーカーであり、これはユーザーの Python プログラムを実行するローカルプロセスとは異なります。そのため、リモートクラスタに接続して TPU を初期化するには、ある程度の初期化作業が必要となります。tf.distribute.cluster_resolver.TPUClusterResolver の tpu 引数は、Colab だけの特別なアドレスであることに注意してください。Google Compute Engine（GCE）で実行している場合は、ご利用の CloudTPU の名前を渡す必要があります。
注意: TPU の初期化コードはプログラムのはじめにある必要があります。
End of explanation
"""


a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

with tf.device('/TPU:0'):
  c = tf.matmul(a, b)

print("c device: ", c.device)
print(c)

"""
Explanation: 手動でデバイスを配置する
TPU が初期されたら、計算を単一の TPU デバイスに配置するために、手動によるデバイスの配置を使用できます。
End of explanation
"""


strategy = tf.distribute.TPUStrategy(resolver)

"""
Explanation: 分散ストラテジー
モデルは通常、複数の TPU で並行して実行されます。複数の TPU（またはその他のアクセラレータ）でモデルを分散させるため、TensorFlow にはいくつかの分散ストラテジーが用意されています。分散ストラテジーを置き換えると、指定された任意の（TPU）デバイスでモデルが実行するようになります。詳細については、分散ストラテジーガイドをご覧ください。
これを実演するために、tf.distribute.TPUStrategy オブジェクトを作成します。
End of explanation
"""


@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)

"""
Explanation: 計算を複製してすべての TPU コアで実行できるようにするには、計算を strategy.run API に渡します。次の例では、すべてのコアが同じ入力 (a, b) を受け入れて、各コアで独立して行列の乗算を実行しています。出力は、すべてのレプリカからの値となります。
End of explanation
"""


def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Conv2D(256, 3, activation='relu'),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256, activation='relu'),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])

"""
Explanation: TPU での分類
基本的な概念を説明したので、より具体的な例を考察しましょう。このセクションでは、分散ストラテジー tf.distribute.TPUStrategy を使用して Cloud TPU でKeras モデルをトレーニングする方法を説明します。
Keras モデルを定義する
MNIST データセットで Keras を使用して画像の分類を行う Sequential Keras モデルの定義から始めましょう。CPU または GPU でトレーニングする場合に使用するものと変わりません。Keras モデルの作成は strategy.scope 内で行う必要があることに注意してください。そうすることで、変数が各 TPU デバイスに作成されるようになります。コードの他の部分は、ストラテジースコープ内にある必要はありません。
End of explanation
"""


def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)

  # Normalize the input data.
  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0
    return image, label

  dataset = dataset.map(scale)

  # Only shuffle and repeat the dataset in training. The advantage of having an
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so that you don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  dataset = dataset.batch(batch_size)

  return dataset

"""
Explanation: データセットを読み込む
Cloud TPU にデータを迅速にフィードできなければ Cloud TPU を使用することは不可能であるため、Cloud TPU を使用する際は、tf.data.Dataset API を効率的に使用できることが非常に重要となります。データセットのパフォーマンスについての詳細は、入力パイプラインのパフォーマンスガイドをご覧ください。
最も単純な実験（tf.data.Dataset.from_tensor_slices またはほかのグラフ内データの使用）以外のすべての実験では、Dataset が読み取るすべてのデータファイルを Google Cloud Storage（GCS）バケットに格納する必要があります。
ほとんどの使用事例では、データを TFRecord 形式に変換し、tf.data.TFRecordDataset を使って読み取ることをお勧めします。このやり方については、「TFRecord および tf.Example のチュートリアル」を参照してください。これは絶対要件ではないため、ほかのデータセットリーダー（tf.data.FixedLengthRecordDataset または tf.data.TextLineDataset）を使用することもできます。
小さなデータセットは、tf.data.Dataset.cache を使ってすべてをメモリに読み込むことができます。
データ形式にかかわらず、100 MB 程度の大きなファイルを使用することをお勧めします。このネットワーク化された設定においては、ファイルを開くタスクのオーバーヘッドが著しく高くなるため、特に重要なことです。
以下のコードに示される通り、tensorflow_datasets モジュールを使用して、MNIST トレーニングデータのコピーを取得する必要があります。try_gcs は、パブリック GCS バケットで提供されているコピーを使用するように指定されています。これを指定しない場合、TPU はダウンロードされたデータにアクセスできません。
End of explanation
"""


with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset, 
          validation_steps=validation_steps)

"""
Explanation: Keras の高位 API を使用してモデルをトレーニングする
Keras の fit と compile API を使用してモデルをトレーニングできます。ここでは、TPU 固有のステップはないため、複数の GPU と MirroredStrategy（TPUStrategy ではなく）を使用している場合と同じようにコードを記述します。詳細については、「Keras を使用した分散トレーニング」チュートリアルをご覧ください。
End of explanation
"""


with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                steps_per_execution = 50,
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset,
          validation_steps=validation_steps)

"""
Explanation: Python のオーバーヘッドを緩和し、TPU のパフォーマンスを最大化するには、引数 steps_per_execution を Model.compile に渡します。この例では、スループットが約 50% 増加します。
End of explanation
"""


# Create the model, optimizer and metrics inside the strategy scope, so that the
# variables can be mirrored on each device.
with strategy.scope():
  model = create_model()
  optimizer = tf.keras.optimizers.Adam()
  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      'training_accuracy', dtype=tf.float32)

# Calculate per replica batch size, and distribute the datasets on each TPU
# worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync

train_dataset = strategy.experimental_distribute_datasets_from_function(
    lambda _: get_dataset(per_replica_batch_size, is_training=True))

@tf.function
def train_step(iterator):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  strategy.run(step_fn, args=(next(iterator),))

"""
Explanation: カスタムトレーニングループを使用してモデルをトレーニングする
tf.function と tf.distribute API を直接使用しても、モデルを作成してトレーニングすることができます。strategy.experimental_distribute_datasets_from_function API は、データセット関数を指定してデータセットを分散させるために使用されます。以下の例では、データセットに渡されるバッチサイズは、グローバルバッチサイズではなく、レプリカごとのバッチサイズであることに注意してください。詳細については、「tf.distribute.Strategy によるカスタムトレーニング」チュートリアルをご覧ください。
最初に、モデル、データセット、および tf.function を作成します。
End of explanation
"""


steps_per_eval = 10000 // batch_size

train_iterator = iter(train_dataset)
for epoch in range(5):
  print('Epoch: {}/5'.format(epoch))

  for step in range(steps_per_epoch):
    train_step(train_iterator)
  print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
  training_loss.reset_states()
  training_accuracy.reset_states()

"""
Explanation: 次に、トレーニングループを実行します。
End of explanation
"""


@tf.function
def train_multiple_steps(iterator, steps):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  for _ in tf.range(steps):
    strategy.run(step_fn, args=(next(iterator),))

# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get 
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))

print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))

"""
Explanation: tf.function 内の複数のステップでパフォーマンスを改善する
tf.function 内で複数のステップを実行することで、パフォーマンスを改善できます。これは、tf.function 内の tf.range で strategy.run 呼び出しをラッピングすることで実現されます。AutoGraph は、TPU ワーカーの tf.while_loop に変換します。
パフォーマンスは改善されますが、tf.function 内の単一のステップに比べれば、この方法にはトレードオフがあります。tf.function で複数のステップを実行すると柔軟性に劣り、ステップ内での Eager execution や任意の Python コードを実行できません。
End of explanation
"""


