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

"""
Explanation: TensorFlow による分散型トレーニング
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://www.tensorflow.org/guide/distributed_training"><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/distributed_training.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/distributed_training.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/distributed_training.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">ノートブックをダウンロード</a></td>
</table>

概要
tf.distribute.Strategy は、複数の GPU、複数のマシン、または TPU でトレーニングを分散する TensorFlow API です。この API を使用すると、最小限のコード変更により、既存のモデルとトレーニングコードを分散することができます。
tf.distribute.Strategy は、次の主な目標を念頭に設計されています。

使いやすく、研究者や ML エンジニアなど、複数のユーザーセグメントをサポートすること。
調整することなく優れたパフォーマンスを提供できること。
ストラテジーを簡単に切り替えられること。

tf.distribute.Strategy は、 Keras Model.fit などの高レベル API とともに使用でき、カスタムトレーニングループを分散するためにも使用できます。
TensorFlow 2.x では、プログラムを Eager モードで、または tf.function を使用してグラフ内で実行することができます。tf.distribute.Strategy はこれらの実行モードをサポートするのが狙いですが、tf.function を使用するのが最適です。Eager モードはデバッグ目的にのみ推奨され、TPUStrategy ではサポートされていません。このガイドの大半はトレーニングに関する内容ですが、異なるプラットフォームに評価と予測を分散する上でも、この API を使用することができます。
tf.distribute.Strategy は、コードをほとんど変更せずに使用することができます。TensorFlow の基盤のコンポーネントをストラテジーを認識するように変更されているからです。これには、変数、レイヤー、モデル、オプティマイザー、メトリック、要約、およびチェックポイントが含まれます。
このガイドでは、さまざまなストラテジーと、異なる状況においてそれらのストラテジーをどのように使用するかについて説明します。
注意: 概念の理解を深めるには、こちらの詳しいプレゼンテーションをご覧ください。独自のトレーニングループを記述する予定の方に、特にお勧めです。
TensorFlow をセットアップする
End of explanation
"""


mirrored_strategy = tf.distribute.MirroredStrategy()

"""
Explanation: ストラテジーの種類
tf.distribute.Strategy は、さまざまな軸に沿って多数の使用事例をカバーすることを目的としています。これらの組み合わせの中には、現在サポートされているものもあれば、今後の追加が予定されているものもあります。次に一部の軸を示します。

同期と非同期トレーニング: これらは、データの並列処理を使用してトレーニングを分散する 2 つの一般的な方法です。同期トレーニングでは、すべてのワーカーは入力データのさまざまなスライスで同期的にトレーニングし、各ステップで勾配を収集します。一方、非同期トレーニングでは、すべてのワーカーは独立的に入力データでトレーニングし、変数を非同期的に更新します。通常、同期トレーニングは all-reduce を介して、非同期トレーニングはパラメーターサーバーアーキテクチャを介してサポートされています。
ハードウェアプラットフォーム: トレーニングを、1台のマシンの複数の GPU またはネットワーク内の複数のマシン（それぞれに 0 個以上の GPU）、さらには Cloud TPU に拡張することができます。

これらのユースケースをサポートするために、TensorFlow では MirroredStrategy、TPUStrategy、MultiWorkerMirroredStrategy、ParameterServerStrategy、CentralStorageStrategy などのストラテジーを利用できます。次のセクションでは、TensorFlow のどのシナリオでこれらのどれがサポートされているかを説明します。以下は簡単な概要です。
トレーニング API | MirroredStrategy | TPUStrategy | MultiWorkerMirroredStrategy | CentralStorageStrategy | ParameterServerStrategy
:-- | :-- | :-- | :-- | :-- | :--
Keras API | サポート中 | サポート中 | 実験的サポート | 実験的サポート | 2.3 以降でサポート予定
カスタムトレーニングループ | サポート中 | サポート中 | 実験的サポート | 実験的サポート | 2.3 以降でサポート予定
Estimator API | 制限サポート | 未サポート | 制限サポート | 制限サポート | 制限サポート
注意: 「実験的サポート」とは、API の互換性が保証されていないことを意味します。
警告: Estimator のサポートは制限されています。基本的なトレーニングと評価は実験的なものであり、スキャフォールドなどの高度な機能は実装されていません。ユースケースがサポートされていない場合は、Keras またはカスタムトレーニングループを使用する必要があります。新しいコードには Estimator は推奨されません。Estimator は v1.Session スタイルのコードを実行しますが、これは正しく記述するのはより難しく、特に TF 2 コードと組み合わせると予期しない動作をする可能性があります。Estimator は、互換性保証の対象となりますが、セキュリティの脆弱性以外の修正は行われません。詳細については、移行ガイドを参照してください。
MirroredStrategy
tf.distribute.MirroredStrategy は、1 台のマシンの複数の GPU での同期分散型トレーニングをサポートしています。GPU デバイス当たり 1 つのレプリカを作成します。モデルの各変数はそのレプリカ全体にミラーリングされます。これらの変数を合わせて形成されるのが、MirroredVariable と呼ばれる 1 つの概念的変数です。これらの変数は、同一の更新を適用することで、相互の同期が維持されます。
デバイス全体に変数の更新を伝達するには、有効性の高い all-reduce アルゴリズムが使用されます。all-reduce はすべてのデバイスのテンソルを加算して集計し、各デバイスで使用できるようにします。非常に効率性が高く、同期のオーバーヘッドを著しく軽減できる融合アルゴリズムです。デバイス間で利用できる通信の種類に応じて、さまざまな all-reduce アルゴリズムと実装が用意されています。デフォルトでは、all-reduce 実装として NVIDIA NCCL が使用されます。提供されているいくつかのオプションから選択できます。また、独自に作成することもできます。
次は、最も単純な MirroredStrategy の作成方法です。
End of explanation
"""


mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

"""
Explanation: これにより、TensorFlow が認識できるすべての GPU を使用する MirroredStrategy インスタンスが作成され、クロスデバイス通信として NCCL が使用されるようになります。
マシン上の一部の GPU のみを使用する場合は、次のように行うことができます。
End of explanation
"""


mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

"""
Explanation: クロスデバイス通信をオーバーライドする場合は、cross_device_ops 引数に tf.distribute.CrossDeviceOps のインスタンスを提供することで実行できます。現在のところ、デフォルトの tf.distribute.NcclAllReduce のほかに、tf.distribute.HierarchicalCopyAllReduce と tf.distribute.ReductionToOneDevice の 2 つのオプションを使用できます。
End of explanation
"""


strategy = tf.distribute.MultiWorkerMirroredStrategy()

"""
Explanation: TPUStrategy
tf.distribute.experimental.TPUStrategy は、Tensor Processing Unit（TPU）で TensorFlow トレーニングを実行できるようにします。TPU は、機械学習ワークロードを劇的に高速化するように設計された Google の特殊 ASIC です。Google Colab の TensorFlow Research Cloud と Cloud TPU で利用できます。
分散型トレーニングアーキテクチャの観点で言えば、TPUStrategy は MirroredStrategy と同じで、同期分散型トレーニングを実装します。TPU は複数の TPU コアに、効率的な all-reduce の独自の実装とほかの集合演算を提供しており、TPUStrategy で使用されています。
以下に、TPUStrategy をインスタンス化する方法を示します。
注意: Colab でこのコードを実行するには、Colab ランタイムとして TPU を選択する必要があります。TensorFlow TPU ガイドを参照してください。
python
cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)
TPUClusterResolver インスタンスを使って TPU の場所を特定することができます。Colab では、引数を指定する必要はありません。
Cloud TPU でこれを使用する場合は、次の内容に注意してくだい。

tpu 引数に、TPU リソースの名前を指定する必要があります。
プログラムの開始時点で、TPU システム明示的に初期化する必要があります。これは、TPU が計算に使用される前に必要なことです。TPU システムを初期化すると、TPU メモリも消去されるため、状態を失わないようにするには、この手順を完了しておくことが重要となります。

MultiWorkerMirroredStrategy
tf.distribute.MultiWorkerMirroredStrategy は、MirroredStrategy に非常に似通ったストラテジーです。複数のワーカーにそれぞれ潜在的に複数の GPU を使って同期分散型トレーニングを実装します。MirroredStrategy と同様に、すべてのワーカーの各デバイスのモデルにすべての変数のコピーを作成します。
次に、MultiWorkerMirroredStrategy を作成する最も簡単な方法を示します。
End of explanation
"""


communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)

"""
Explanation: MultiWorkerMirroredStrategy では、現在、2 つの collective 演算の実装を選択することができます。 CollectiveCommunication.RING は、gRPC を使用してリング型の collective を通信レイヤーとして実装します。 CollectiveCommunication.NCCL は、Nvidia の NCCL を使用して集合を実装します。CollectiveCommunication.AUTO は、選択をランタイムに任せます。collective 実装の最適な選択は、GPU の数と種類、およびクラスタ内のネットワーク相互接続によって決まります。これらは、次のようにして指定することができます。
End of explanation
"""


central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

"""
Explanation: マルチワーカートレーニングを進める上でマルチ GPU トレーニングとの主な違いの 1 つに、マルチワーカーのセットアップが挙げられます。TF_CONFIG 環境変数は、TensorFlow においてクラスタ構成をそのクラスタの一部である各ワーカーに指定する標準的な方法です。TF_CONFIG の設定についてさらにお読みください。
MultiWorkerMirroredStrategy についての詳細は、以下のチュートリアルを参照してくだい。

Keras Model.fit を使ったマルチワーカートレーニング
カスタムトレーニングループを使ったマルチワーカートレーニング

ParameterServerStrategy
パラメータサーバートレーニングは、複数のマシンでモデルトレーニングをスケールアップするための一般的なデータ並列方法です。パラメータサーバートレーニングクラスタは、ワーカーとパラメータサーバーで構成されます。変数はパラメータサーバー上に作成され、各ステップでワーカーにより読み取られ、更新されます。詳細については、パラメータサーバートレーニングチュートリアルを参照してください。
TensorFlow 2 では、パラメータサーバートレーニングは、tf.distribute.experimental.coordinator.ClusterCoordinator クラスを介して中央コーディネーターベースのアーキテクチャを使用します。
この実装では、worker と parameter server タスクは、コーディネーターからのタスクをリッスンする tf.distribute.Serverを実行します。コーディネーターは、リソースを作成し、トレーニングタスクをディスパッチし、チェックポイントを書き込み、タスクの失敗に対処します。
コーディネーターで実行されているプログラミングでは、ParameterServerStrategy オブジェクトを使用してトレーニングステップを定義し、ClusterCoordinator を使用してトレーニングステップをリモートワーカーにディスパッチします。これらを作成する最も簡単な方法は次のとおりです。
python
strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)
ParameterServerStrategy の詳細については、 Keras Model.fit を使用したパラメータサーバートレーニングとカスタムトレーニングループチュートリアルを参照してください。
注意: TFConfigClusterResolver を使用する場合は、TF_CONFIG 環境変数を構成する必要があります。<a href="#TF_CONFIG" data-md-type="link">TF_CONFIG</a>（MultiWorkerMirroredStrategy）に似ていますが、別の注意点があります。
TensorFlow 1 では、ParameterServerStrategy{ code0}は、&lt;code data-md-type="codespan"&gt;tf.compat.v1.distribute.experimental.ParameterServerStrategyシンボルを介した Estimator でのみ使用できます。
注意: このストラテジーは現在改善中であり、より多くのシナリオで機能できるようにしていることから、「実験的」に指定されています。この改善プロセスの一環として、API の振る舞いが今後変更される可能性があることにご注意ください。
CentralStorageStrategy
tf.distribute.experimental.CentralStorageStrategy も同期トレーニングを行います。変数はミラーリングされず、代わりに CPU に配置され、演算はすべてのローカル GPU に複製されます。GPU が 1 つしかない場合は、すべての変数と演算はその GPU に配置されます。
次のようにして、CentralStorageStrategy のインスタンスを作成します。
End of explanation
"""


default_strategy = tf.distribute.get_strategy()

"""
Explanation: これにより、すべての GPU と CPU を使用する CentralStorageStrategy インスタンスが作成されます。レプリカでの変数の更新は、変数に適用される前に集計されます。
注意: このストラテジーは現在改善中であり、より多くのシナリオで機能できるようにしていることから、「実験的」に指定されています。この改善プロセスの一環として、API の振る舞いが今後変更される可能性があることにご注意ください。
その他のストラテジー
上述のストラテジーのほかに、tf.distribute API を使用する際のプロトタイピングとデバッグに役立つ可能性のあるストラテジーが 2 つあります。
デフォルトストラテジー
明示的な分散ストラテジーがスコープにない場合には、分散ストラテジーとしてデフォルトストラテジーが使用されます。tf.distribute.Strategy インターフェースを実装しますが、パススルーであるため、実際の分散を提供しません。たとえば、strategy.run(fn) は fn だけを呼び出します。このストラテジーを使用して書かれたコードは、ストラテジーを指定せずに書かれたコードとまった同じ振る舞いとなります。「演算なし」ストラテジーとして考えるとわかりやすいかもしれません。
デフォルトストラテジーはシングルトンであり、インスタンスを 1 つしか作成できません。明示的なストラテジーのスコープの外で、tf.distribute.get_strategy() を使って取得することができます（明示的なストラテジーのスコープ内で現在のストラテジーを取得するために使用するのと同じ API）。
End of explanation
"""


# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values

"""
Explanation: このストラテジーには、主に 2 つの目的があります。

無条件で分散対応のライブラリコードを記述すること。たとえば、tf.optimizer において、tf.distribute.get_strategy を実行し、勾配を減らすためにそのストラテジーを使用します。これは常に、Strategy.reduce API を呼び出せるストラテジーオブジェクトを返します。
End of explanation
"""


if tf.config.list_physical_devices('GPU'):
  strategy = tf.distribute.MirroredStrategy()
else:  # Use the Default Strategy
  strategy = tf.distribute.get_strategy()

with strategy.scope():
  # Do something interesting
  print(tf.Variable(1.))

"""
Explanation: ライブラリコードと同様に、条件論理を必要とせずに、分散ストラテジーの有無に関係なく機能するエンドユーザーのプログラムを記述すること。これを説明したサンプルコードスニペットを、次に示します。
End of explanation
"""


mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')

"""
Explanation: OneDeviceStrategy
tf.distribute.OneDeviceStrategy は、すべての変数と計算を単一の指定デバイスに配置するストラテジーです。
python
strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
このストラテジーは、多数の点において、デフォルトストラテジーと異なります。デフォルトストラテジーでは、分散ストラテジーを使用せずに TensorFlow を実行した場合と比べ、変数の配置ロジックに変化はありません。しかし、OneDeviceStrategy を使用した場合、そのスコープで作成されるすべての変数は、指定されたデバイスに明示的に配置されます。さらに、OneDeviceStrategy.run から呼び出される関数も、その指定デバイスに配置されるようになります。
このストラテジーを通じて分散された入力は、指定デバイスにプリフェッチされます。デフォルトストラテジーには、入力の分散はありません。
デフォルトストラテジーと同様に、このストラテジーも、複数のデバイス/マシンに実際に分散する別のストラテジーに切り替える前のコードテストに使用することができます。これにより、分散ストラテジーの仕組みはある程度はデフォルトストラテジーよりもが強化されますが、MirroredStrategy や TPUStrategy などを使用した場合ほど最大限には強化されません。ストラテジーが指定されていないかのようなコードの振る舞いを希望する場合は、デフォルトストラテジーを使用してください。
以上、利用できるさまざまなストラテジーとそのインスタンス化の方法を見てきました。以下では、トレーニングを分散化するために使用できるさまざまな方法を見ていきます。
tf.distribute.Strategy を Keras Model.fit と使用する
TensorFlow の Keras API 仕様の実装である tf.keras に tf.distribute.Strategy を統合しました。tf.keras は、モデルを構築してトレーニングするための高位 API です。tf.keras バックエンドに統合することによって、model.fit を使用する Keras トレーニングフレームワークに記述されたトレーニングの分散をシームレスに行えるようになっています。
コードを次のように変更してください。

適切な tf.distribute.Strategy のインスタンスを作成します。
Keras モデル、オプティマイザー、およびメトリックの作成を strategy.scope 内に移動します。

TensorFlow 分散ストラテジーは、Sequential、Functional、および Subclassed のすべての Keras モデルをサポートしています。
次に、1 つの Dense レイヤーを持つ非常に単純な Keras モデルに対してこれを実行するコードスニペットを示します。
End of explanation
"""


dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)

"""
Explanation: この例では、複数の GPU を持つマシンで実行できるように、MirroredStrategy を使用しています。strategy.scope() は、トレーニングの分散に度のストラテジーを使用するかを Keras に示しています。このスコープ内にモデル/オプティマイザー/メトリックを作成することで、通常の変数ではなく、分散化された変数を作成することができます。これをセットアップしたら、通常通り、モデルをフィットさせることができます。利用できる GPU に対するモデルのトレーニングの複製や勾配の集計などは、MirroredStrategy によって行われます。
End of explanation
"""


import numpy as np

inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)

"""
Explanation: ここでは、tf.data.Dataset を使用してトレーニングと eval 入力を提供していますが、次のように numpy 配列を使用することもできます。
End of explanation
"""


# Compute a global batch size using a number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

"""
Explanation: いずれの場合（Dataset または Numpy）でも、指定された入力の各バッチは、複数のレプリカで均等に分割されます。たとえば、2 GPU で MirroredStrategy を使用している場合、サイズが 10 の各バッチは 2 GPU 間で分割されるため、各ステップでそれぞれに 5 つの入力例が提供されます。その後、GPU をさらに追加するにつれ、各エポックはより高速にトレーニングするようになります。一般的に、アクセラレータを増やすたびにバッチサイズを増加することが望ましく、そうすることで、追加の計算パワーを有効活用できるようになります。また、モデルに応じて学習速度を再調整することも必要です。レプリカの数を取得するには、strategy.num_replicas_in_sync を使用できます。
End of explanation
"""


with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  optimizer = tf.keras.optimizers.SGD()

"""
Explanation: 現在、何がサポートされていますか？
トレーニング API | MirroredStrategy | TPUStrategy | MultiWorkerMirroredStrategy | ParameterServerStrategy | CentralStorageStrategy
--- | --- | --- | --- | --- | ---
Keras API | サポート中 | サポート中 | 実験的サポート | 実験的サポート | 2.3 以降でサポート予定
例とチュートリアル
次は、上述した Keras のエンドツーエンド統合を説明するチュートリアルと例の一覧です。

チュートリアル: Model.fit と MirroredStrategy でトレーニングする。
チュートリアル: Model.fit と MultiWorkerMirroredStrategy でトレーニングする。
ガイド:  Model.fit と TPUStrategy を使用する例を含むガイド。
チュートリアル: Model.fit と ParameterServerStrategy でパラメータサーバーをトレーニングする。
チュートリアル: Model.fit と TPUStrategy を使用して、GLUE ベンチマークから多くのタスクの BERT を微調整する。
さまざまなストラテジーを使って実装された最新モデルのコレクションが含まれる TensorFlow Model Garden リポジトリ

カスタムトレーニングループで tf.distribute.Strategy <br> を使用する
これまで見てきたように、Keras の model.fit で tf.distribute.Strategy を使用するにはコードの数行のみを変更する必要がありました。もう少し手を加えれば、カスタムトレーニングループでも tf.distribute.Strategy を使用できるようになります。
トレーニングループに Estimator や Keras よりも高い柔軟性と制御が必要な場合は、カスタムトレーニングループを作成することができます。たとえば、GANを使用する際に、各ラウンドで異なる数のジェネレータやディスクリミネータのステップを実行することができます。同様に、高度なフレームワークは、強化学習トレーニングにはあまり適していません。
カスタムとレニングループをサポートするために、tf.distribute.Strategy クラスでメソッドのコアセットを提供しています。これらを使用するには、最初にコードをわずかに再構成する必要があるかもしれませんが、それを一度行うと、ストラテジーインスタンスを変更するだけで、GPU、TPU、および複数のマシンを切り替えられるようになります。
ここでは、以前と同じ Keras モデルを使って、単純なトレーニングでの使用事例を表した簡単なスニペットを示します。
まず、ストラテジーのスコープ内にモデルとオプティマイザーを作成します。こうすることで、そのモデルとオプティマイザーを使って作成された変数がミラーリングされた変数であることを保証することができます。
End of explanation
"""


dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(
    global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

"""
Explanation: 次に、入力データセットを作成し、tf.distribute.Strategy.experimental_distribute_dataset を呼び出して、ストラテジーに応じてデータセットを分散します。
End of explanation
"""


loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def compute_loss(labels, predictions):
  per_example_loss = loss_object(labels, predictions)
  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    loss = compute_loss(labels, predictions)

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

@tf.function
def distributed_train_step(dist_inputs):
  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

"""
Explanation: そして、トレーニングの一ステップを定義します。勾配とオプティマイザーを計算し、その勾配を適用してモデルの変数を更新するために、tf.GradientTape を使用します。このトレーニングステップを分散するには、train_step 関数に入れて、前に作成した dist_dataset から取得するデータセット入力とともに、tf.distrbute.Strategy.run に渡します。
End of explanation
"""


for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))

"""
Explanation: 上記のコードには、注意すべき点がいくつかあります。

tf.nn.compute_average_loss を使用して損失を計算しました。tf.nn.compute_average_loss は例ごとの損失を加算し、その和を global_batch_size で除算します。これは、各レプリカで勾配が計算された後、その勾配を加算することで、レプリカ全体の勾配が集計されるため重要です。
tf.distribute.Strategy.run が返す結果を集計するために、tf.distribute.Strategy.reduce API を使用しました。tf.distribute.Strategy.run は、ストラテジーの各ローカルレプリカの結果を返し、この結果の消費方法もさまざまです。集計された値を取得するには、結果を reduce することができます。また、tf.distribute.Strategy.experimental_local_results を実行して、結果に含まれる、ローカルレプリカ当たり 1 つの値のリストを取得することもできます。
分散ストラテジーのスコープ内で apply_gradients が呼び出されると、その振る舞いが変更されます。具体的には、同期トレーニング中に各並列インスタンスに勾配を適用する前に、勾配の sum-over-all-replicas を実行します。

最後に、トレーニングステップの定義が済んだら、dist_dataset をイテレートしてトレーニングをループで実行できます。
End of explanation
"""


iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))

"""
Explanation: 上記の例では、dist_dataset をイテレートして、トレーニングに入力を提供しました。また、numpy 入力をサポートするために、tf.distribute.Strategy.make_experimental_numpy_dataset も提供しています。tf.distribute.Strategy.experimental_distribute_dataset を呼び出す前にこの API を使って、データセットを作成することができます。
データをイテレートするための別の方法は、明示的にイテレータを使用することです。データセット全体をイテレートするのに対し、特定のステップ数だけ実行する場合に、これを行うことができます。上記のイテレートは、最初にイテレータを作成し、それに明示的にnext を呼び出して入力データを取得するように変更されます。
End of explanation
"""


