 TensorFlow.org で表示 |
 Google Colab で実行 |
 GitHub でソースを表示 |
 ノートブックをダウンロード |
概要
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 GPU パフォーマンスを最適化する」ガイドをご覧ください。
注意: 概念の理解を深めるには、こちらの詳しいプレゼンテーションをご覧ください。独自のトレーニングループを記述する予定の方に、特にお勧めです。
TensorFlow をセットアップする
import tensorflow as tf
2024-01-11 20:33:52.811223: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2024-01-11 20:33:52.811273: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2024-01-11 20:33:52.813068: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
ストラテジーの種類
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 Model.fit |
サポート中 | サポート中 | 実験的サポート | 実験的サポート | 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 の作成方法です。
mirrored_strategy = tf.distribute.MirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
これにより、TensorFlow が認識できるすべての GPU を使用する MirroredStrategy インスタンスが作成され、クロスデバイス通信として NCCL が使用されるようになります。
マシン上の一部の GPU のみを使用する場合は、次のように行うことができます。
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')
クロスデバイス通信をオーバーライドする場合は、cross_device_ops 引数に tf.distribute.CrossDeviceOps のインスタンスを提供することで実行できます。現在のところ、デフォルトの tf.distribute.NcclAllReduce のほかに、tf.distribute.HierarchicalCopyAllReduce と tf.distribute.ReductionToOneDevice の 2 つのオプションを使用できます。
mirrored_strategy = tf.distribute.MirroredStrategy(
cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
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 ガイドを参照してください。
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.experimental.TPUStrategy(cluster_resolver)
TPUClusterResolver インスタンスを使って TPU の場所を特定することができます。Colab では、引数を指定する必要はありません。
Cloud TPU でこれを使用する場合は、次の内容に注意してくだい。
tpu引数に、TPU リソースの名前を指定する必要があります。- プログラムの開始時点で、TPU システム明示的に初期化する必要があります。これは、TPU が計算に使用される前に必要なことです。TPU システムを初期化すると、TPU メモリも消去されるため、状態を失わないようにするには、この手順を完了しておくことが重要となります。
MultiWorkerMirroredStrategy
tf.distribute.experimental.MultiWorkerMirroredStrategy は、MirroredStrategy に非常に似通ったストラテジーです。複数のワーカーにそれぞれ潜在的に複数の GPU を使って同期分散型トレーニングを実装します。MirroredStrategy と同様に、すべてのワーカーの各デバイスのモデルにすべての変数のコピーを作成します。
マルチワーカー all-reduce 通信方法として CollectiveOps を使用して、変数の同期を維持しています。collective 演算は、TensorFlow グラフの単一の演算で、ハードウェア、ネットワークテクノロジー、およびテンソルのサイズに応じて TensorFlow ランタイムに all-reduce アルゴリズムを自動的に選択することができます。
strategy = tf.distribute.MultiWorkerMirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0', '/device:GPU:1', '/device:GPU:2', '/device:GPU:3')
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0', '/device:GPU:1', '/device:GPU:2', '/device:GPU:3'), communication = CommunicationImplementation.AUTO
MultiWorkerMirroredStrategy では、現在、2 つの collective 演算の実装を選択することができます。 CollectiveCommunication.RING は、gRPC を使用してリング型の collective を通信レイヤーとして実装します。 CollectiveCommunication.NCCL は、Nvidia の NCCL を使用して集合を実装します。CollectiveCommunication.AUTO は、選択をランタイムに任せます。collective 実装の最適な選択は、GPU の数と種類、およびクラスタ内のネットワーク相互接続によって決まります。これらは、次のようにして指定することができます。
communication_options = tf.distribute.experimental.CommunicationOptions(
implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
communication_options=communication_options)
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0', '/device:GPU:1', '/device:GPU:2', '/device:GPU:3')
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0', '/device:GPU:1', '/device:GPU:2', '/device:GPU:3'), communication = CommunicationImplementation.NCCL
マルチワーカートレーニングを進める上でマルチ GPU トレーニングとの主な違いの 1 つに、マルチワーカーのセットアップが挙げられます。TF_CONFIG 環境変数は、TensorFlow においてクラスタ構成をそのクラスタの一部である各ワーカーに指定する標準的な方法です。TF_CONFIG の設定についてさらにお読みください。
MultiWorkerMirroredStrategy についての詳細は、以下のチュートリアルを参照してくだい。
ParameterServerStrategy
パラメータサーバートレーニングは、複数のマシンでモデルトレーニングをスケールアップするための一般的なデータ並列方法です。パラメータサーバートレーニングクラスタは、ワーカーとパラメータサーバーで構成されます。変数はパラメータサーバー上に作成され、各ステップでワーカーにより読み取られ、更新されます。詳細については、パラメータサーバートレーニングチュートリアルを参照してください。
TensorFlow 2 では、パラメータサーバートレーニングは、tf.distribute.experimental.coordinator.ClusterCoordinator クラスを介して中央コーディネーターベースのアーキテクチャを使用します。
この実装では、worker と parameter server タスクは、コーディネーターからのタスクをリッスンする tf.distribute.Server を実行します。コーディネーターは、リソースを作成し、トレーニングタスクをディスパッチし、チェックポイントを書き込み、タスクの失敗に対処します。
コーディネーターで実行されているプログラミングでは、ParameterServerStrategy オブジェクトを使用してトレーニングステップを定義し、ClusterCoordinator を使用してトレーニングステップをリモートワーカーにディスパッチします。これらを作成する最も簡単な方法は次のとおりです。
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' 環境変数を構成する必要があります。MultiWorkerMirroredStrategy の 'TF_CONFIG' に似ていますが、別の注意点があります。
TensorFlow 1 では、ParameterServerStrategy{ code0} は、<code data-md-type="codespan">tf.compat.v1.distribute.experimental.ParameterServerStrategy シンボルを介した Estimator でのみ使用できます。
注意: このストラテジーは現在改善中であり、より多くのシナリオで機能できるようにしていることから、「実験的」に指定されています。この改善プロセスの一環として、API の振る舞いが今後変更される可能性があることにご注意ください。
CentralStorageStrategy
tf.distribute.experimental.CentralStorageStrategy も同期トレーニングを行います。変数はミラーリングされず、代わりに CPU に配置され、演算はすべてのローカル GPU に複製されます。GPU が 1 つしかない場合は、すべての変数と演算はその GPU に配置されます。
次のようにして、CentralStorageStrategy のインスタンスを作成します。
central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3'], variable_device = '/device:CPU:0'
これにより、すべての GPU と CPU を使用する CentralStorageStrategy インスタンスが作成されます。レプリカでの変数の更新は、変数に適用される前に集計されます。
注意: このストラテジーは進行中の作業であり、experimental となっています。
その他のストラテジー
上述のストラテジーのほかに、tf.distribute API を使用する際のプロトタイピングとデバッグに役立つ可能性のあるストラテジーが 2 つあります。
デフォルトストラテジー
明示的な分散ストラテジーがスコープにない場合には、分散ストラテジーとしてデフォルトストラテジーが使用されます。tf.distribute.Strategy インターフェースを実装しますが、パススルーであるため、実際の分散を提供しません。たとえば、strategy.run(fn) は fn だけを呼び出します。このストラテジーを使用して書かれたコードは、ストラテジーを指定せずに書かれたコードとまった同じ振る舞いとなります。「演算なし」ストラテジーとして考えるとわかりやすいかもしれません。
デフォルトストラテジーはシングルトンであり、インスタンスを 1 つしか作成できません。明示的なストラテジーのスコープの外で、tf.distribute.get_strategy() を使って取得することができます（明示的なストラテジーのスコープ内で現在のストラテジーを取得するために使用するのと同じ API）。
default_strategy = tf.distribute.get_strategy()
このストラテジーには、主に 2 つの目的があります。
- 無条件で分散対応のライブラリコードを記述すること。たとえば、
tf.keras.optimizersにおいて、tf.distribute.get_strategyを実行し、勾配を減らすためにそのストラテジーを使用します。これは常に、Strategy.reduceAPI を呼び出せるストラテジーオブジェクトを返します。
# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None) # reduce some values
1.0
- ライブラリコードと同様に、条件論理を必要とせずに、分散ストラテジーの有無に関係なく機能するエンドユーザーのプログラムを記述すること。これを説明したサンプルコードスニペットを、次に示します。
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.))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
MirroredVariable:{
0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,
1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=1.0>,
2: <tf.Variable 'Variable/replica_2:0' shape=() dtype=float32, numpy=1.0>,
3: <tf.Variable 'Variable/replica_3:0' shape=() dtype=float32, numpy=1.0>
}
OneDeviceStrategy
tf.distribute.OneDeviceStrategy は、すべての変数と計算を単一の指定デバイスに配置するストラテジーです。
strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
このストラテジーは、多数の点において、デフォルトストラテジーと異なります。デフォルトストラテジーでは、分散ストラテジーを使用せずに TensorFlow を実行した場合と比べ、変数の配置ロジックに変化はありません。しかし、OneDeviceStrategy を使用した場合、そのスコープで作成されるすべての変数は、指定されたデバイスに明示的に配置されます。さらに、OneDeviceStrategy.run から呼び出される関数も、その指定デバイスに配置されるようになります。
このストラテジーを通じて分散された入力は、指定デバイスにプリフェッチされます。デフォルトストラテジーには、入力の分散はありません。
デフォルトストラテジーと同様に、このストラテジーも、複数のデバイス/マシンに実際に分散する別のストラテジーに切り替える前のコードテストに使用することができます。これにより、分散ストラテジーの仕組みはある程度はデフォルトストラテジーよりもが強化されますが、MirroredStrategy や TPUStrategy などを使用した場合ほど最大限には強化されません。ストラテジーが指定されていないかのようなコードの振る舞いを希望する場合は、デフォルトストラテジーを使用してください。
以上、利用できるさまざまなストラテジーとそのインスタンス化の方法を見てきました。以下では、トレーニングを分散化するために使用できるさまざまな方法を見ていきます。
tf.keras.Model.fit を使った tf.distribute.Strategy の使用
tf.distribute.Strategy は、TensorFlow の Keras API 仕様の実装である tf.keras に統合されています。tf.keras は、モデルを構築してトレーニングするための高位 API です。tf.keras バックエンドに統合することによって、Model.fitを使用する Keras トレーニングフレームワークに記述されたトレーニングの分散をシームレスに行えるようになっています。
コードを次のように変更してください。
- 適切な
tf.distribute.Strategyのインスタンスを作成します。 - 作成した Keras モデル、オプティマイザ、および指標を
strategy.scopeに移動します。モデルのcall()、train_step()、およびtest_step()メソッド内のコードは、分散化されてアクセラレータで実行されます。
TensorFlow 分散ストラテジーは、Sequential、Functional、および subclassed のすべての Keras モデルをサポートしています。
次に、1 つの Dense レイヤーを持つ非常に単純な Keras モデルに対してこれを実行するコードスニペットを示します。
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
model = tf.keras.Sequential([
tf.keras.layers.Dense(1, input_shape=(1,),
kernel_regularizer=tf.keras.regularizers.L2(1e-4))])
model.compile(loss='mse', optimizer='sgd')
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
この例では、複数の GPU を持つマシンで実行できるように、MirroredStrategy を使用しています。strategy.scope() は、トレーニングの分散に度のストラテジーを使用するかを Keras に示しています。このスコープ内にモデル/オプティマイザー/メトリックを作成することで、通常の変数ではなく、分散化された変数を作成することができます。これをセットアップしたら、通常通り、モデルをフィットさせることができます。利用できる GPU に対するモデルのトレーニングの複製や勾配の集計などは、MirroredStrategy によって行われます。
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)
Epoch 1/2
INFO:tensorflow:Collective all_reduce tensors: 2 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
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:Collective all_reduce tensors: 2 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
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',).
1/10 [==>...........................] - ETA: 31s - loss: 1.0575e-04INFO: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',).
10/10 [==============================] - 4s 7ms/step - loss: 1.0416e-04
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1705005242.237586 283178 device_compiler.h:186] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process.
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',).
Epoch 2/2
1/10 [==>...........................] - ETA: 0s - loss: 1.0267e-04INFO: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',).
10/10 [==============================] - 0s 6ms/step - loss: 1.0196e-04
10/10 [==============================] - 1s 4ms/step - loss: 1.0130e-04
0.0001013019063975662
ここでは、tf.data.Dataset を使用してトレーニングと eval 入力を提供していますが、次のように NumPy 配列を使用することもできます。
import numpy as np
inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)
Epoch 1/2 10/10 [==============================] - 0s 5ms/step - loss: 1.0099e-04 Epoch 2/2 10/10 [==============================] - 0s 5ms/step - loss: 1.0056e-04 <keras.src.callbacks.History at 0x7f4670335940>
いずれの場合（Dataset または Numpy）でも、指定された入力の各バッチは、複数のレプリカで均等に分割されます。たとえば、2 GPU で MirroredStrategy を使用している場合、サイズが 10 の各バッチは 2 GPU 間で分割されるため、各ステップでそれぞれに 5 つの入力例が提供されます。その後、GPU をさらに追加するにつれ、各エポックはより高速にトレーニングするようになります。一般的に、アクセラレータを増やすたびにバッチサイズを増加することが望ましく、そうすることで、追加の計算パワーを有効活用できるようになります。また、モデルに応じて学習速度を再調整することも必要です。レプリカの数を取得するには、strategy.num_replicas_in_sync を使用できます。
mirrored_strategy.num_replicas_in_sync
4
# 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, 20:0.175}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]
現在、何がサポートされていますか？
| トレーニング API | MirroredStrategy |
TPUStrategy |
MultiWorkerMirroredStrategy |
ParameterServerStrategy |
CentralStorageStrategy |
|---|---|---|---|---|---|
| Keras API | サポート中 | サポート中 | サポート中 | 実験的サポート | 2.3 以降でサポート予定 |
例とチュートリアル
次は、上述した Keras のエンドツーエンド統合を説明するチュートリアルと例の一覧です。
MirroredStrategyで MNIST をトレーニングするチュートリアルMultiWorkerMirroredStrategyを使って MNIST をトレーニングするチュートリアル- ガイド:
Model.fitとTPUStrategyを使用する例を含むガイド。 - チュートリアル:
Model.fitとParameterServerStrategyでパラメータサーバーをトレーニングする。 - チュートリアル:
Model.fitとTPUStrategyを使用して、GLUE ベンチマークから多くのタスクの BERT を微調整する。 - さまざまなストラテジーを使って実装された最新モデルのコレクションが含まれる TensorFlow Model Garden リポジトリ
カスタムトレーニングループで tf.distribute.Strategy
を使用する
これまで見てきたように、Keras の Model.fit で tf.distribute.Strategy を使用するにはコードの数行のみを変更する必要がありました。もう少し手を加えれば、カスタムトレーニングループでも tf.distribute.Strategy を使用できるようになります。
トレーニングループに Estimator や Keras よりも高い柔軟性と制御が必要な場合は、カスタムトレーニングループを作成することができます。たとえば、GANを使用する際に、各ラウンドで異なる数のジェネレータやディスクリミネータのステップを実行することができます。同様に、高度なフレームワークは、強化学習トレーニングにはあまり適していません。
カスタムトレニングループをサポートするために、tf.distribute.Strategy クラスでメソッドのコアセットを提供しています。これらを使用するには、最初にコードをわずかに再構成する必要があるかもしれませんが、それを一度行うと、ストラテジーインスタンスを変更するだけで、GPU、TPU、および複数のマシンを切り替えられるようになります。
ここでは、以前と同じ Keras モデルを使って、単純なトレーニングでの使用事例を表した簡単なスニペットを示します。
まず、ストラテジーのスコープ内にモデルとオプティマイザーを作成します。こうすることで、そのモデルとオプティマイザーを使って作成された変数がミラーリングされた変数であることを保証することができます。
with mirrored_strategy.scope():
model = tf.keras.Sequential([
tf.keras.layers.Dense(1, input_shape=(1,),
kernel_regularizer=tf.keras.regularizers.L2(1e-4))])
optimizer = tf.keras.optimizers.SGD()
次に、入力データセットを作成し、tf.distribute.Strategy.experimental_distribute_dataset を呼び出して、ストラテジーに応じてデータセットを分散します。
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
そして、トレーニングの一ステップを定義します。勾配とオプティマイザーを計算し、その勾配を適用してモデルの変数を更新するために、tf.GradientTape を使用します。このトレーニングステップを分散するには、train_step 関数に入れて、前に作成した dist_dataset から取得するデータセット入力とともに、tf.distrbute.Strategy.run に渡します。
# Sets `reduction=NONE` to leave it to tf.nn.compute_average_loss() below.
loss_object = tf.keras.losses.BinaryCrossentropy(
from_logits=True,
reduction=tf.keras.losses.Reduction.NONE)
def train_step(inputs):
features, labels = inputs
with tf.GradientTape() as tape:
predictions = model(features, training=True)
per_example_loss = loss_object(labels, predictions)
loss = tf.nn.compute_average_loss(per_example_loss)
model_losses = model.losses
if model_losses:
loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))
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)
上記のコードには、注意すべき点がいくつかあります。
tf.nn.compute_average_lossを使用して、サンプルごとの予測損失をスカラーに縮小しました。tf.nn.compute_average_lossはサンプルごとの損失を加算し、その和をグローバルバッチサイズで除算しています。これは、各レプリカで勾配が計算された後、その勾配を加算することで、レプリカ全体の勾配が集計されるため重要です。
デフォルトでは、tf.get_strategy().num_replicas_in_sync * tf.shape(per_example_loss)[0] になるようにグローバルバッチサイズが取られます。これはキーワード引数 global_batch_size= として明示的に指定することも可能です。短いバッチがない場合、このデフォルトは、tf.nn.compute_average_loss(..., global_batch_size=global_batch_size) と上記で定義した global_batch_size と同等です。（短いバッチ、およびそれらの回避または処理の方法については、カスタムトレーニングチュートリアルをご覧ください。）
tf.nn.scale_regularization_lossを使用して、Modelオブジェクトに登録されている正則化損失（存在する場合）も、1/num_replicas_in_syncによってスケーリングしました。入力に依存する正則化損失の場合、カスタムトレーニングループではなくモデリングコードでレプリカ(!) ごとのバッチサイズの平均化を実行します。そうすることで、モデリングコードはレプリケーションに、トレーニングループは正則化損失の計算方法に依存することがありません。分散ストラテジーのスコープ内で
apply_gradientsが呼び出されると、その振る舞いが変更されます。具体的には、同期トレーニング中に各並列インスタンスに勾配を適用する前に、勾配の sum-over-all-replicas を実行します。tf.distribute.Strategy.runが返した結果を集計するために、tf.distribute.Strategy.reduceAPI も使用しました。tf.distribute.Strategy.runはストラテジー内の各ローカルレプリカの結果を返し、この結果の使用方法は多様です。reduceで、集約された値を取得することができます。また、tf.distribute.Strategy.experimental_local_resultsを実行して、ローカルレプリカごとに 1 つ、結果に含まれる値のリストを取得することもできます。
最後に、トレーニングステップの定義が済んだら、dist_dataset をイテレートしてトレーニングをループで実行できます。
for dist_inputs in dist_dataset:
print(distributed_train_step(dist_inputs))
INFO:tensorflow:Collective all_reduce tensors: 2 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1 INFO:tensorflow:Collective all_reduce tensors: 2 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1 tf.Tensor(1.0082505, shape=(), dtype=float32) tf.Tensor(1.0002, shape=(), dtype=float32) tf.Tensor(0.9922241, shape=(), dtype=float32) tf.Tensor(0.9843224, shape=(), dtype=float32) tf.Tensor(0.97649455, shape=(), dtype=float32) tf.Tensor(0.96874, shape=(), dtype=float32) tf.Tensor(0.96105826, shape=(), dtype=float32) tf.Tensor(0.9534487, shape=(), dtype=float32) tf.Tensor(0.9459112, shape=(), dtype=float32) tf.Tensor(0.9384451, shape=(), dtype=float32) tf.Tensor(0.9310499, shape=(), dtype=float32) tf.Tensor(0.92372507, shape=(), dtype=float32) tf.Tensor(0.91647035, shape=(), dtype=float32) tf.Tensor(0.909285, shape=(), dtype=float32) tf.Tensor(0.90216863, shape=(), dtype=float32) tf.Tensor(0.8951206, shape=(), dtype=float32) tf.Tensor(0.8881406, shape=(), dtype=float32) tf.Tensor(0.8812279, shape=(), dtype=float32) tf.Tensor(0.8743823, shape=(), dtype=float32) tf.Tensor(0.867603, shape=(), dtype=float32) tf.Tensor(0.8608897, shape=(), dtype=float32) tf.Tensor(0.8542416, shape=(), dtype=float32) tf.Tensor(0.8476584, shape=(), dtype=float32) tf.Tensor(0.8411393, shape=(), dtype=float32) tf.Tensor(0.8346843, shape=(), dtype=float32) tf.Tensor(0.82829237, shape=(), dtype=float32) tf.Tensor(0.8219633, shape=(), dtype=float32) tf.Tensor(0.8156963, shape=(), dtype=float32) tf.Tensor(0.80949104, shape=(), dtype=float32) tf.Tensor(0.80334705, shape=(), dtype=float32) tf.Tensor(0.79726344, shape=(), dtype=float32) tf.Tensor(0.79124004, shape=(), dtype=float32) tf.Tensor(0.7852762, shape=(), dtype=float32) tf.Tensor(0.7793713, shape=(), dtype=float32) tf.Tensor(0.773525, shape=(), dtype=float32) tf.Tensor(0.7677367, shape=(), dtype=float32) tf.Tensor(0.7620057, shape=(), dtype=float32) tf.Tensor(0.75633174, shape=(), dtype=float32) tf.Tensor(0.75071406, shape=(), dtype=float32) tf.Tensor(0.74515235, shape=(), dtype=float32) tf.Tensor(0.739646, shape=(), dtype=float32) tf.Tensor(0.73419434, shape=(), dtype=float32) tf.Tensor(0.7287971, shape=(), dtype=float32) tf.Tensor(0.7234536, shape=(), dtype=float32) tf.Tensor(0.7181633, shape=(), dtype=float32) tf.Tensor(0.71292585, shape=(), dtype=float32) tf.Tensor(0.70774055, shape=(), dtype=float32) tf.Tensor(0.7026069, shape=(), dtype=float32) tf.Tensor(0.6975246, shape=(), dtype=float32) tf.Tensor(0.69249296, shape=(), dtype=float32)
上記の例では、dist_dataset をイテレートして、トレーニングに入力を提供しました。また、numpy 入力をサポートするために、tf.distribute.Strategy.make_experimental_numpy_dataset も提供しています。tf.distribute.Strategy.experimental_distribute_dataset を呼び出す前にこの API を使って、データセットを作成することができます。
データをイテレートするための別の方法は、明示的にイテレータを使用することです。データセット全体をイテレートするのに対し、特定のステップ数だけ実行する場合に、これを行うことができます。上記のイテレートは、最初にイテレータを作成し、それに明示的にnext を呼び出して入力データを取得するように変更されます。
iterator = iter(dist_dataset)
for _ in range(10):
print(distributed_train_step(next(iterator)))
tf.Tensor(0.6875115, shape=(), dtype=float32) tf.Tensor(0.6825797, shape=(), dtype=float32) tf.Tensor(0.6776972, shape=(), dtype=float32) tf.Tensor(0.6728633, shape=(), dtype=float32) tf.Tensor(0.6680777, shape=(), dtype=float32) tf.Tensor(0.6633398, shape=(), dtype=float32) tf.Tensor(0.6586491, shape=(), dtype=float32) tf.Tensor(0.65400517, shape=(), dtype=float32) tf.Tensor(0.6494075, shape=(), dtype=float32) tf.Tensor(0.6448556, shape=(), dtype=float32)
これは、tf.distribute.Strategy APIを使用してカスタムトレーニングループを分散する最も単純な事例をカバーしています。これらの API は改善プロセスにあります。この使用事例にはコードを適合させる作業がさらに必要であるため、いずれ、詳細なガイドを別途公開する予定です。
現在、何がサポートされていますか？
| トレーニング API | MirroredStrategy |
TPUStrategy |
MultiWorkerMirroredStrategy |
ParameterServerStrategy |
CentralStorageStrategy |
|---|---|---|---|---|---|
| カスタムトレーニングループ | サポート中 | サポート中 | サポート中 | 実験的サポート | 2.3 以降でサポート予定 |
例とチュートリアル
次は、カスタムトレーニングループを使って分散ストラテジーを使用するいくつかの例です。
MirroredStrategyで MNIST をトレーニングするチュートリアルTPUStrategyによる MNIST のトレーニングに関するガイド- さまざまなストラテジーを使って実装された最新モデルのコレクションが含まれる TensorFlow Model Garden リポジトリ
- チュートリアル: カスタムトレーニングループと
ParameterServerStrategyでパラメータサーバーをトレーニングする。 - さまざまなストラテジーを使って実装された最新モデルのコレクションが含まれる TensorFlow Model Garden リポジトリ
その他のトピック
このセクションでは、複数の使用事例に関連のあるトピックをいくつかカバーしています。
TF_CONFIG 環境変数のセットアップ
マルチワーカートレーニングでは、前述のとおり、クラスタで実行されている各バイナリに対して TF_CONFIG 環境変数を設定する必要があります。TF_CONFIG 環境変数は、クラスタを構成するタスク、そのアドレス、およびクラスタ内の各タスクのロールを指定する JSON 文字列です。tensorflow/ecosystem レポジトリには、トレーニングタスクの TF_CONFIG を設定する Kubernotes テンプレートを提供しています。
TF_CONFIG には、cluster と task の 2 つのコンポーネントがあります。
- cluster はトレーニングクラスタに関する情報を提供します。これは、worker などのさまざまなタイプのジョブで構成されるディクショナリーです。マルチワーカートレーニングでは、通常、一般的なワーカーの作業に加えて、チェックポイントの保存や TensorBoard のサマリーファイルの書き込みなど、ほかよりタスクを担うワーカーが 1 つあります。こういったワーカーは、「チーフ」ワーカーと呼ばれ、index 0 のワーカーがチーフワーカーに指定されるようになっています（実際、このようにして、
tf.distribute.Strategyが実装されています）。 - 一方、task は現在のタスクに関する情報を提供します。最初のコンポーネント cluster はすべてのワーカーで同じであり、2番目のコンポーネント task はワーカーごとに異なり、そのワーカーのタイプとインデックスを指定します。
次は、TF_CONFIG の一例です。
os.environ["TF_CONFIG"] = json.dumps({ "cluster": { "worker": ["host1:port", "host2:port", "host3:port"], "ps": ["host4:port", "host5:port"] }, "task": {"type": "worker", "index": 1} })
この 'TF_CONFIG'は、"cluster" 内に 3 つのワーカーと 2 つの "ps" タスク、およびそれらのホストとポートがあることを指定しています。 "task" の部分は、"cluster" の現在のタスクのロールである worker 1（2番目のワーカー）を指定します。クラスタ内の有効なロールは、"chief"、"worker"、"ps"、および "evaluator" です。tf.distribute.experimental.ParameterServerStrategyを使用する場合を除いて、"ps" ジョブはありません。
次のステップ
tf.distribute.Strategy は積極的な開発が進められています。ぜひお試しいただき、GitHub 課題 に皆さんの感想をお寄せください。