![]() | ![]() | ![]() | ![]() |
概要
パラメータサーバートレーニングは、複数のマシンでモデルトレーニングをスケールアップするための一般的なデータ並列方法です。
パラメーターサーバートレーニングクラスターは、ワーカーとパラメーターサーバーで構成されます。変数はパラメーターサーバー上に作成され、各ステップでワーカーによって読み取られて更新されます。デフォルトでは、ワーカーはこれらの変数を相互に同期せずに個別に読み取り、更新します。これが、パラメータサーバースタイルのトレーニングが非同期トレーニングと呼ばれることがある理由です。
TensorFlow 2では、パラメーターサーバーのトレーニングはtf.distribute.experimental.ParameterServerStrategy
クラスによって強化されます。このクラスは、トレーニングステップを数千人のワーカー(パラメーターサーバーを伴う)にスケールアップするクラスターに分散します。
サポートされているトレーニング方法
サポートされている主なトレーニング方法は2つあります。
- Keras
Model.fit
APIは、高レベルの抽象化とトレーニングの処理を希望する場合に推奨されます。 - カスタムトレーニングループ(詳細については、カスタムトレーニング、ゼロからのトレーニングループの作成、およびKerasとMultiWorkerMirroredStrategyを使用したカスタムトレーニングループを参照できます)。トレーニングループの詳細を定義する場合は、カスタムループトレーニングをお勧めします。
ジョブとタスクを含むクラスター
選択したAPI( Model.fit
またはカスタムトレーニングループ)に関係なく、TensorFlow 2の分散トレーニングには、複数'jobs'
を含む'cluster'
が含まれ、各ジョブには1つ以上'tasks'
が含まれる場合があります。
パラメータサーバートレーニングを使用する場合は、次のことをお勧めします。
- 1つのコーディネータージョブ(ジョブ名は
chief
) - 複数のワーカージョブ(ジョブ名
worker
);と - 複数パラメータサーバージョブ(ジョブ名
ps
)
コーディネーターがリソースを作成し、トレーニングタスクをディスパッチし、チェックポイントを書き込み、タスクの失敗に対処している間、ワーカーとパラメーターサーバーはコーディネーターからの要求をリッスンするtf.distribute.Server
を実行します。
Model.fit
を使用したパラメーターサーバーのトレーニング
Model.fit
APIを使用したパラメーターサーバーのトレーニングでは、コーディネーターがtf.distribute.experimental.ParameterServerStrategy
オブジェクトとtf.keras.utils.experimental.DatasetCreator
を入力として使用する必要があります。ストラテジーがない、または他のストラテジーがあるModel.fit
の使用法と同様に、ワークフローには、モデルの作成とコンパイル、コールバックの準備、それに続くModel.fit
呼び出しが含まれます。
カスタムトレーニングループを使用したパラメータサーバートレーニング
カスタムトレーニングループでは、 tf.distribute.experimental.coordinator.ClusterCoordinator
クラスがコーディネーターに使用される主要なコンポーネントです。
-
ClusterCoordinator
クラスは、tf.distribute.Strategy
オブジェクトと連携して機能する必要があります。 - この
tf.distribute.Strategy
オブジェクトは、クラスターの情報を提供するために必要であり、tf.distribute.Strategyを使用したカスタムトレーニングで示されているように、トレーニングステップを定義するために使用されます。 - 次に、
ClusterCoordinator
オブジェクトは、これらのトレーニングステップの実行をリモートワーカーにディスパッチします。 - パラメーターサーバーのトレーニングでは、
ClusterCoordinator
はtf.distribute.experimental.ParameterServerStrategy
と連携する必要があります。
ClusterCoordinator
オブジェクトによって提供される最も重要なAPIはschedule
です:
-
schedule
APIはtf.function
をキューに入れ、futureのようなRemoteValue
をすぐに返します。 - キューに入れられた関数は、バックグラウンドスレッドのリモートワーカーにディスパッチされ、それらの
RemoteValue
は非同期で埋められます。 -
schedule
はワーカーの割り当てを必要としないため、渡されたtf.function
は、使用可能な任意のワーカーで実行できます。 - 実行されたワーカーが完了する前に使用できなくなった場合、関数は別の使用可能なワーカーで再試行されます。
- この事実と関数の実行がアトミックではないという事実のために、関数は複数回実行される可能性があります。
ClusterCoordinator
は、リモート関数のディスパッチに加えて、すべてのワーカーでデータセットを作成し、ワーカーが障害から回復したときにこれらのデータセットを再構築するのにも役立ちます。
チュートリアルのセットアップ
チュートリアルはModel.fit
とカスタムトレーニングループパスに分岐し、ニーズに合ったものを選択できます。 「Xを使用したトレーニング」以外のセクションは、両方のパスに適用されます。
pip install portpicker
import multiprocessing
import os
import random
import portpicker
import tensorflow as tf
クラスターのセットアップ
上記のように、パラメーターサーバートレーニングクラスターには、トレーニングプログラムを実行するコーディネータータスク、TensorFlowサーバーを実行する1つまたは複数のワーカーとパラメーターサーバータスク( tf.distribute.Server
)、および場合によってはサイドカー評価を実行する追加の評価タスクが必要です。 (以下のサイドカー評価セクションを参照してください)。それらを設定するための要件は次のとおりです。
- コーディネータータスクは、エバリュエーターを除く他のすべてのTensorFlowサーバーのアドレスとポートを知る必要があります。
- ワーカーとパラメーターサーバーは、リッスンする必要のあるポートを知る必要があります。簡単にするために、通常、これらのタスクでTensorFlowサーバーを作成するときに、完全なクラスター情報を渡すことができます。
- 評価者タスクは、トレーニングクラスターの設定を知る必要はありません。含まれている場合は、トレーニングクラスターへの接続を試行しないでください。
- ワーカーとパラメーターサーバーは、それぞれ
"worker"
と"ps"
のタスクタイプを持つ必要があります。コーディネーターは、従来の理由から、タスクタイプとして"chief"
を使用する必要があります。
このチュートリアルでは、パラメーターサーバーのトレーニング全体をColabで実行できるように、インプロセスクラスターを作成します。後のセクションで、実際のクラスターを設定する方法を学習します。
インプロセスクラスター
まず、事前に複数のTensorFlowサーバーを作成し、後でそれらに接続します。これはこのチュートリアルのデモンストレーションのみを目的としていることに注意してください。実際のトレーニングでは、サーバーは"worker"
および"ps"
マシンで起動されます。
def create_in_process_cluster(num_workers, num_ps):
"""Creates and starts local servers and returns the cluster_resolver."""
worker_ports = [portpicker.pick_unused_port() for _ in range(num_workers)]
ps_ports = [portpicker.pick_unused_port() for _ in range(num_ps)]
cluster_dict = {}
cluster_dict["worker"] = ["localhost:%s" % port for port in worker_ports]
if num_ps > 0:
cluster_dict["ps"] = ["localhost:%s" % port for port in ps_ports]
cluster_spec = tf.train.ClusterSpec(cluster_dict)
# Workers need some inter_ops threads to work properly.
worker_config = tf.compat.v1.ConfigProto()
if multiprocessing.cpu_count() < num_workers + 1:
worker_config.inter_op_parallelism_threads = num_workers + 1
for i in range(num_workers):
tf.distribute.Server(
cluster_spec,
job_name="worker",
task_index=i,
config=worker_config,
protocol="grpc")
for i in range(num_ps):
tf.distribute.Server(
cluster_spec,
job_name="ps",
task_index=i,
protocol="grpc")
cluster_resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(
cluster_spec, rpc_layer="grpc")
return cluster_resolver
# Set the environment variable to allow reporting worker and ps failure to the
# coordinator. This is a workaround and won't be necessary in the future.
os.environ["GRPC_FAIL_FAST"] = "use_caller"
NUM_WORKERS = 3
NUM_PS = 2
cluster_resolver = create_in_process_cluster(NUM_WORKERS, NUM_PS)
インプロセスクラスターセットアップは、ここなどの単体テストで頻繁に使用されます。
ローカルテストのもう1つのオプションは、ローカルマシンでプロセスを起動することです。このアプローチの例については、 Kerasを使用したマルチワーカートレーニングを確認してください。
ParameterServerStrategyをインスタンス化します
トレーニングコードに飛び込む前に、 ParameterServerStrategy
オブジェクトをインスタンス化しましょう。これは、 Model.fit
を続行するか、カスタムトレーニングループを続行するかに関係なく必要であることに注意してください。 variable_partitioner
引数については、「変数シャーディング」セクションで説明します。
variable_partitioner = (
tf.distribute.experimental.partitioners.MinSizePartitioner(
min_shard_bytes=(256 << 10),
max_shards=NUM_PS))
strategy = tf.distribute.experimental.ParameterServerStrategy(
cluster_resolver,
variable_partitioner=variable_partitioner)
INFO:tensorflow:`tf.distribute.experimental.ParameterServerStrategy` is initialized with cluster_spec: ClusterSpec({'ps': ['localhost:19069', 'localhost:19342'], 'worker': ['localhost:18007', 'localhost:20252', 'localhost:23613']}) INFO:tensorflow:ParameterServerStrategyV2 is now connecting to cluster with cluster_spec: ClusterSpec({'ps': ['localhost:19069', 'localhost:19342'], 'worker': ['localhost:18007', 'localhost:20252', 'localhost:23613']}) INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:chief/replica:0/task:0/device:GPU:0'], variable_device = '/job:chief/replica:0/task:0/device:GPU:0' INFO:tensorflow:Number of GPUs on workers: 1
トレーニングにGPUを使用するには、各ワーカーに表示されるGPUを割り当てます。 ParameterServerStrategy
は、各ワーカーで使用可能なすべてのGPUを使用しますが、すべてのワーカーが同じ数のGPUを使用できるようにする必要があるという制限があります。
可変シャーディング
変数シャーディングとは、変数を複数の小さな変数に分割することを指します。これはシャードと呼ばれます。可変シャーディングは、これらのシャードにアクセスするときにネットワーク負荷を分散するのに役立つ場合があります。正規変数の計算と保存を複数のパラメーターサーバーに分散することも役立ちます。
変数のシャーディングを有効にするには、 ParameterServerStrategy
オブジェクトを作成するときにvariable_partitioner
を渡すことができます。 variable_partitioner
は、変数が作成されるたびに呼び出され、変数の各次元に沿ったシャードの数を返すことが期待されます。 tf.distribute.experimental.partitioners.MinSizePartitioner
など、すぐに使用できるvariable_partitioner
がいくつか提供されています。モデルのトレーニング速度に悪影響を与える可能性のある小さな変数の分割を回避するために、 tf.distribute.experimental.partitioners.MinSizePartitioner
などのサイズベースのパーティショナーを使用することをお勧めします。
variable_partitioner
が渡され、 strategy.scope()
の直下で変数を作成すると、シャードのリストへのアクセスを提供するvariables
プロパティを持つコンテナータイプになります。ほとんどの場合、このコンテナはすべてのシャードを連結することで自動的にTensorに変換されます。その結果、正規変数として使用できます。一方、 tf.nn.embedding_lookup
などの一部のTensorFlowメソッドは、このコンテナタイプの効率的な実装を提供し、これらのメソッドでは自動連結が回避されます。
詳細については、 tf.distribute.experimental.ParameterServerStrategy
のAPIドキュメントを参照してください。
Model.fit
を使用したトレーニング
Kerasは、 Model.fit
を介して使いやすいトレーニングAPIを提供します。このAPIは、オーバーライド可能なtrain_step
の柔軟性を備えた内部でトレーニングループを処理し、コールバックはTensorBoardのチェックポイント保存やサマリー保存などの機能を提供します。 Model.fit
を使用すると、戦略オブジェクトを簡単に交換するだけで、同じトレーニングコードを他の戦略に使用できます。
入力データ
パラメータサーバートレーニングを使用したModel.fit
では、入力データが、タイプtf.distribute.InputContext
の単一の引数を取り、 tf.data.Dataset
を返す呼び出し可能オブジェクトで提供される必要があります。次に、そのようなcallable
オブジェクトを受け取るtf.distribute.InputOptions
オブジェクトと、 input_options
引数を介したオプションのtf.keras.utils.experimental.DatasetCreator
オブジェクトを作成します。
パラメータサーバートレーニングを使用してデータをシャッフルして繰り返し、ライブラリがエポック境界を認識できるように、 fit
呼び出しでsteps_per_epoch
を指定することをお勧めします。
InputContext
引数の詳細については、分散入力チュートリアルを参照してください。
def dataset_fn(input_context):
global_batch_size = 64
batch_size = input_context.get_per_replica_batch_size(global_batch_size)
x = tf.random.uniform((10, 10))
y = tf.random.uniform((10,))
dataset = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(10).repeat()
dataset = dataset.shard(
input_context.num_input_pipelines,
input_context.input_pipeline_id)
dataset = dataset.batch(batch_size)
dataset = dataset.prefetch(2)
return dataset
dc = tf.keras.utils.experimental.DatasetCreator(dataset_fn)
dataset_fn
のコードは、各ワーカーマシンの入力デバイス(通常はCPU)で呼び出されます。
モデルの構築とコンパイル
ここで、 tf.keras.models.Sequential
tf.keras.Model
)を作成し、続いてModel.compile
呼び出しを作成して、オプティマイザー、メトリック、 steps_per_execution
などのパラメーターなどのコンポーネントを組み込みます。
with strategy.scope():
model = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
model.compile(tf.keras.optimizers.SGD(), loss='mse', steps_per_execution=10)
コールバックとトレーニング
実際のトレーニングのためにmodel.fit
を呼び出す前に、次のような一般的なタスクに必要なコールバックを準備しましょう。
-
ModelCheckpoint
:モデルの重みを保存します。 -
BackupAndRestore
:トレーニングの進行状況が自動的にバックアップされ、クラスターで使用不可(中止やプリエンプションなど)が発生した場合に回復されるようにします。また TensorBoard
:進捗レポートを要約ファイルに保存します。要約ファイルはTensorBoardツールで視覚化されます。
working_dir = '/tmp/my_working_dir'
log_dir = os.path.join(working_dir, 'log')
ckpt_filepath = os.path.join(working_dir, 'ckpt')
backup_dir = os.path.join(working_dir, 'backup')
callbacks = [
tf.keras.callbacks.TensorBoard(log_dir=log_dir),
tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_filepath),
tf.keras.callbacks.BackupAndRestore(backup_dir=backup_dir),
]
model.fit(dc, epochs=5, steps_per_epoch=20, callbacks=callbacks)
Epoch 1/5 INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). 2022-01-26 05:32:01.399347: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them. INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets 20/20 - 4s - loss: 0.5761 - 4s/epoch - 185ms/step Epoch 2/5 INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets 20/20 - 1s - loss: 0.4423 - 561ms/epoch - 28ms/step Epoch 3/5 WARNING:tensorflow:5 out of the last 5 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7f89783e7560> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details. INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets WARNING:tensorflow:6 out of the last 6 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7f897851f050> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details. 20/20 - 0s - loss: 0.3983 - 392ms/epoch - 20ms/step Epoch 4/5 INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets 20/20 - 1s - loss: 0.3592 - 507ms/epoch - 25ms/step Epoch 5/5 INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets 20/20 - 0s - loss: 0.2705 - 384ms/epoch - 19ms/step <keras.callbacks.History at 0x7f89984ca890>
ClusterCoordinator
での直接使用(オプション)
Model.fit
トレーニングパスを選択した場合でも、オプションでtf.distribute.experimental.coordinator.ClusterCoordinator
オブジェクトをインスタンス化して、ワーカーで実行する他の関数をスケジュールできます。詳細と例については、カスタムトレーニングループを使用したトレーニングのセクションを参照してください。
カスタムトレーニングループを使用したトレーニング
tf.distribute.Strategy
でカスタムトレーニングループを使用すると、トレーニングループを定義するための優れた柔軟性が得られます。上で( strategy
として)定義されたParameterServerStrategy
を使用して、 tf.distribute.experimental.coordinator.ClusterCoordinator
を使用して、トレーニングステップの実行をリモートワーカーにディスパッチします。
次に、他のtf.distribute.Strategy
を使用したトレーニングループで行ったように、モデルを作成し、データセットとステップ関数を定義します。詳細については、 tf.distribute.Strategyチュートリアルを使用したカスタムトレーニングを参照してください。
効率的なデータセットのプリフェッチを確実にするには、以下のリモートワーカーへのディスパッチトレーニング手順で説明されている推奨される分散データセット作成APIを使用してください。また、 worker_fn
内でStrategy.run
を呼び出して、ワーカーに割り当てられたGPUを最大限に活用してください。残りの手順は、GPUを使用する場合と使用しない場合のトレーニングで同じです。
次の手順でこれらのコンポーネントを作成しましょう。
データを設定する
まず、 Keras前処理レイヤーによって実装された前処理ロジックを含むデータセットを作成する関数を記述します。
これらのレイヤーはdataset_fn
の外部に作成しますが、 dataset_fn
をtf.function
にラップするため、 dataset_fn
の内部に変換を適用します。これにより、変数を内部に作成できなくなります。
feature_vocab = [
"avenger", "ironman", "batman", "hulk", "spiderman", "kingkong", "wonder_woman"
]
label_vocab = ["yes", "no"]
with strategy.scope():
feature_lookup_layer = tf.keras.layers.StringLookup(
vocabulary=feature_vocab,
mask_token=None)
label_lookup_layer = tf.keras.layers.StringLookup(
vocabulary=label_vocab,
num_oov_indices=0,
mask_token=None)
raw_feature_input = tf.keras.layers.Input(
shape=(3,),
dtype=tf.string,
name="feature")
feature_id_input = feature_lookup_layer(raw_feature_input)
feature_preprocess_stage = tf.keras.Model(
{"features": raw_feature_input},
feature_id_input)
raw_label_input = tf.keras.layers.Input(
shape=(1,),
dtype=tf.string,
name="label")
label_id_input = label_lookup_layer(raw_label_input)
label_preprocess_stage = tf.keras.Model(
{"label": raw_label_input},
label_id_input)
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/numpy/core/numeric.py:2446: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison return bool(asarray(a1 == a2).all())
データセットにおもちゃの例を生成します。
def feature_and_label_gen(num_examples=200):
examples = {"features": [], "label": []}
for _ in range(num_examples):
features = random.sample(feature_vocab, 3)
label = ["yes"] if "avenger" in features else ["no"]
examples["features"].append(features)
examples["label"].append(label)
return examples
examples = feature_and_label_gen()
次に、 dataset_fn
でラップされたトレーニングデータセットを作成します。
def dataset_fn(_):
raw_dataset = tf.data.Dataset.from_tensor_slices(examples)
train_dataset = raw_dataset.map(
lambda x: (
{"features": feature_preprocess_stage(x["features"])},
label_preprocess_stage(x["label"])
)).shuffle(200).batch(32).repeat()
return train_dataset
モデルを構築する
次に、モデルとその他のオブジェクトを作成します。必ずstrategy.scope
の下にすべての変数を作成してください。
# These variables created under the `strategy.scope` will be placed on parameter
# servers in a round-robin fashion.
with strategy.scope():
# Create the model. The input needs to be compatible with Keras processing layers.
model_input = tf.keras.layers.Input(
shape=(3,), dtype=tf.int64, name="model_input")
emb_layer = tf.keras.layers.Embedding(
input_dim=len(feature_lookup_layer.get_vocabulary()), output_dim=16384)
emb_output = tf.reduce_mean(emb_layer(model_input), axis=1)
dense_output = tf.keras.layers.Dense(units=1, activation="sigmoid")(emb_output)
model = tf.keras.Model({"features": model_input}, dense_output)
optimizer = tf.keras.optimizers.RMSprop(learning_rate=0.1)
accuracy = tf.keras.metrics.Accuracy()
FixedShardsPartitioner
を使用すると、すべての変数が2つのシャードに分割され、各シャードが異なるパラメーターサーバーに割り当てられていることを確認しましょう。
assert len(emb_layer.weights) == 2
assert emb_layer.weights[0].shape == (4, 16384)
assert emb_layer.weights[1].shape == (4, 16384)
assert emb_layer.weights[0].device == "/job:ps/replica:0/task:0/device:CPU:0"
assert emb_layer.weights[1].device == "/job:ps/replica:0/task:1/device:CPU:0"
トレーニングステップを定義する
第三に、 tf.function
にラップされたトレーニングステップを作成します。
@tf.function
def step_fn(iterator):
def replica_fn(batch_data, labels):
with tf.GradientTape() as tape:
pred = model(batch_data, training=True)
per_example_loss = tf.keras.losses.BinaryCrossentropy(
reduction=tf.keras.losses.Reduction.NONE)(labels, pred)
loss = tf.nn.compute_average_loss(per_example_loss)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
accuracy.update_state(labels, actual_pred)
return loss
batch_data, labels = next(iterator)
losses = strategy.run(replica_fn, args=(batch_data, labels))
return strategy.reduce(tf.distribute.ReduceOp.SUM, losses, axis=None)
上記のトレーニングステップ関数では、 step_fn
でStrategy.run
とStrategy.reduce
を呼び出すと、ワーカーごとに複数のGPUをサポートできます。ワーカーにGPUが割り当てられている場合、 Strategy.run
はデータセットを複数のレプリカに分散します。
リモートワーカーへのトレーニング手順のディスパッチ
すべての計算がParameterServerStrategy
によって定義されたら、 tf.distribute.experimental.coordinator.ClusterCoordinator
クラスを使用してリソースを作成し、トレーニング手順をリモートワーカーに配布します。
まず、 ClusterCoordinator
オブジェクトを作成し、戦略オブジェクトを渡します。
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(strategy)
次に、ワーカーごとのデータセットとイテレーターを作成します。以下のper_worker_dataset_fn
では、GPUへの効率的なプリフェッチをシームレスに行えるように、 dataset_fn
をstrategy.distribute_datasets_from_function
にラップすることをお勧めします。
@tf.function
def per_worker_dataset_fn():
return strategy.distribute_datasets_from_function(dataset_fn)
per_worker_dataset = coordinator.create_per_worker_dataset(per_worker_dataset_fn)
per_worker_iterator = iter(per_worker_dataset)
WARNING:tensorflow:Model was constructed with shape (None, 3) for input KerasTensor(type_spec=TensorSpec(shape=(None, 3), dtype=tf.string, name='feature'), name='feature', description="created by layer 'feature'"), but it was called on an input with incompatible shape (3,).プレースホルダー19
最後のステップは、 ClusterCoordinator.schedule
を使用してリモートワーカーに計算を分散することです。
-
schedule
メソッドはtf.function
をキューに入れ、futureのようなRemoteValue
をすぐに返します。キューに入れられた関数はバックグラウンドスレッドのリモートワーカーにディスパッチされ、RemoteValue
は非同期で埋められます。 -
join
メソッド(ClusterCoordinator.join
)を使用して、スケジュールされたすべての関数が実行されるまで待機できます。
num_epoches = 4
steps_per_epoch = 5
for i in range(num_epoches):
accuracy.reset_states()
for _ in range(steps_per_epoch):
coordinator.schedule(step_fn, args=(per_worker_iterator,))
# Wait at epoch boundaries.
coordinator.join()
print ("Finished epoch %d, accuracy is %f." % (i, accuracy.result().numpy()))
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',). Finished epoch 0, accuracy is 0.543750. Finished epoch 1, accuracy is 0.543750. Finished epoch 2, accuracy is 0.950000. Finished epoch 3, accuracy is 1.000000.
RemoteValue
の結果を取得する方法は次のとおりです。
loss = coordinator.schedule(step_fn, args=(per_worker_iterator,))
print ("Final loss is %f" % loss.fetch())
Final loss is 0.000000
または、すべてのステップを起動して、完了を待っている間に何かを実行することもできます。
for _ in range(total_steps):
coordinator.schedule(step_fn, args=(per_worker_iterator,))
while not coordinator.done():
time.sleep(10)
# Do something like logging metrics or writing checkpoints.
この特定の例の完全なトレーニングとサービスのワークフローについては、このテストを確認してください。
データセット作成の詳細
上記のコードのデータセットは、 ClusterCoordinator.create_per_worker_dataset
を使用して作成されています。ワーカーごとに1つのデータセットを作成し、コンテナーオブジェクトを返します。 iter
メソッドを呼び出して、ワーカーごとのイテレーターを作成できます。ワーカーごとのイテレーターには、ワーカーごとに1つのイテレーターが含まれ、特定のワーカーで関数が実行される前に、 ClusterCoordinator.schedule
メソッドに渡される関数の入力引数でワーカーの対応するスライスが置き換えられます。
現在、 ClusterCoordinator.schedule
メソッドは、ワーカーが同等であると想定しているため、 Dataset.shuffle
操作が含まれている場合に異なる方法でシャッフルされる可能性があることを除いて、異なるワーカーのデータセットは同じであると想定しています。このため、データセットを無期限に繰り返し、データセットのOutOfRangeError
に依存するのではなく、有限数のステップをスケジュールすることもお勧めします。
もう1つの重要な注意点は、 tf.data
データセットは、タスクの境界を越えた暗黙的なシリアル化と逆シリアル化をサポートしていないことです。したがって、 ClusterCoordinator.create_per_worker_dataset
に渡される関数内にデータセット全体を作成することが重要です。
評価
分散トレーニングで評価ループを定義して実行する方法は複数あります。以下に説明するように、それぞれに独自の長所と短所があります。好みがない場合は、インライン評価方法をお勧めします。
インライン評価
この方法では、コーディネーターがトレーニングと評価を交互に行うため、インライン評価と呼ばれます。
インライン評価にはいくつかの利点があります。例えば:
- 単一のタスクでは保持できない大規模な評価モデルと評価データセットをサポートできます。
- 評価結果は、次のエポックをトレーニングするための決定を行うために使用できます。
インライン評価を実装するには、直接評価と分散評価の2つの方法があります。
- 直接評価:小さなモデルと評価データセットの場合、コーディネーターは、コーディネーターの評価データセットを使用して、分散モデルで直接評価を実行できます。
eval_dataset = tf.data.Dataset.from_tensor_slices(
feature_and_label_gen(num_examples=16)).map(
lambda x: (
{"features": feature_preprocess_stage(x["features"])},
label_preprocess_stage(x["label"])
)).batch(8)
eval_accuracy = tf.keras.metrics.Accuracy()
for batch_data, labels in eval_dataset:
pred = model(batch_data, training=False)
actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
eval_accuracy.update_state(labels, actual_pred)
print ("Evaluation accuracy: %f" % eval_accuracy.result())
WARNING:tensorflow:Model was constructed with shape (None, 3) for input KerasTensor(type_spec=TensorSpec(shape=(None, 3), dtype=tf.string, name='feature'), name='feature', description="created by layer 'feature'"), but it was called on an input with incompatible shape (3,). Evaluation accuracy: 1.000000プレースホルダー26
- 分散評価:コーディネーターで直接実行することが不可能な大規模なモデルまたはデータセットの場合、コーディネータータスクは
ClusterCoordinator.schedule
ClusterCoordinator.join
メソッドを介して評価タスクをワーカーに分散できます。
with strategy.scope():
# Define the eval metric on parameter servers.
eval_accuracy = tf.keras.metrics.Accuracy()
@tf.function
def eval_step(iterator):
def replica_fn(batch_data, labels):
pred = model(batch_data, training=False)
actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
eval_accuracy.update_state(labels, actual_pred)
batch_data, labels = next(iterator)
strategy.run(replica_fn, args=(batch_data, labels))
def eval_dataset_fn():
return tf.data.Dataset.from_tensor_slices(
feature_and_label_gen(num_examples=16)).map(
lambda x: (
{"features": feature_preprocess_stage(x["features"])},
label_preprocess_stage(x["label"])
)).shuffle(16).repeat().batch(8)
per_worker_eval_dataset = coordinator.create_per_worker_dataset(eval_dataset_fn)
per_worker_eval_iterator = iter(per_worker_eval_dataset)
eval_steps_per_epoch = 2
for _ in range(eval_steps_per_epoch):
coordinator.schedule(eval_step, args=(per_worker_eval_iterator,))
coordinator.join()
print ("Evaluation accuracy: %f" % eval_accuracy.result())
WARNING:tensorflow:Model was constructed with shape (None, 3) for input KerasTensor(type_spec=TensorSpec(shape=(None, 3), dtype=tf.string, name='feature'), name='feature', description="created by layer 'feature'"), but it was called on an input with incompatible shape (3,). WARNING:tensorflow:1 GPUs are allocated per worker. Please use DistributedDataset by calling strategy.experimental_distribute_dataset or strategy.distribute_datasets_from_function to make best use of GPU resources WARNING:tensorflow:1 GPUs are allocated per worker. Please use DistributedDataset by calling strategy.experimental_distribute_dataset or strategy.distribute_datasets_from_function to make best use of GPU resources WARNING:tensorflow:1 GPUs are allocated per worker. Please use DistributedDataset by calling strategy.experimental_distribute_dataset or strategy.distribute_datasets_from_function to make best use of GPU resources Evaluation accuracy: 1.000000プレースホルダー28
サイドカー評価
もう1つの方法は、サイドカー評価と呼ばれ、チェックポイントを繰り返し読み取り、最新のチェックポイントで評価を実行する専用の評価タスクを作成します。評価結果に基づいてトレーニングループを変更する必要がない場合は、トレーニングプログラムを早期に終了できます。ただし、評価をトリガーするには、追加の評価タスクと定期的なチェックポイントが必要です。以下は、可能なサイドカー評価ループです。
checkpoint_dir = ...
eval_model = ...
eval_data = ...
checkpoint = tf.train.Checkpoint(model=eval_model)
for latest_checkpoint in tf.train.checkpoints_iterator(
checkpoint_dir):
try:
checkpoint.restore(latest_checkpoint).expect_partial()
except (tf.errors.OpError,) as e:
# checkpoint may be deleted by training when it is about to read it.
continue
# Optionally add callbacks to write summaries.
eval_model.evaluate(eval_data)
# Evaluation finishes when it has evaluated the last epoch.
if latest_checkpoint.endswith('-{}'.format(train_epoches)):
break
実世界のクラスター
実際の実稼働環境では、すべてのタスクをさまざまなマシンのさまざまなプロセスで実行します。各タスクでクラスター情報を構成する最も簡単な方法は、 "TF_CONFIG"
環境変数を設定し、 tf.distribute.cluster_resolver.TFConfigClusterResolver
を使用して"TF_CONFIG"
を解析することです。
"TF_CONFIG"
環境変数の一般的な説明については、分散トレーニングガイドを参照してください。
Kubernetesまたはその他の構成テンプレートを使用してトレーニングタスクを開始する場合、これらのテンプレートですでに“TF_CONFIG"
が設定されている可能性があります。
"TF_CONFIG"
環境変数を設定します
3つのワーカーと2つのパラメーターサーバーがあるとすると、ワーカー1の"TF_CONFIG"
は次のようになります。
os.environ["TF_CONFIG"] = json.dumps({
"cluster": {
"worker": ["host1:port", "host2:port", "host3:port"],
"ps": ["host4:port", "host5:port"],
"chief": ["host6:port"]
},
"task": {"type": "worker", "index": 1}
})
評価者の"TF_CONFIG"
は次のようになります。
os.environ["TF_CONFIG"] = json.dumps({
"cluster": {
"evaluator": ["host7:port"]
},
"task": {"type": "evaluator", "index": 0}
})
上記のエバリュエーターの"TF_CONFIG"
文字列の"cluster"
部分はオプションです。
すべてのタスクに同じバイナリを使用する場合
単一のバイナリを使用してこれらすべてのタスクを実行する場合は、最初にプログラムをさまざまな役割に分岐させる必要があります。
cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()
if cluster_resolver.task_type in ("worker", "ps"):
# Start a TensorFlow server and wait.
elif cluster_resolver.task_type == "evaluator":
# Run side-car evaluation
else:
# Run the coordinator.
次のコードはTensorFlowサーバーを起動し、待機します。
# Set the environment variable to allow reporting worker and ps failure to the
# coordinator. This is a workaround and won't be necessary in the future.
os.environ["GRPC_FAIL_FAST"] = "use_caller"
server = tf.distribute.Server(
cluster_resolver.cluster_spec(),
job_name=cluster_resolver.task_type,
task_index=cluster_resolver.task_id,
protocol=cluster_resolver.rpc_layer or "grpc",
start=True)
server.join()
タスクの失敗の処理
労働者の失敗
tf.distribute.experimental.coordinator.ClusterCoordinator
またはModel.fit
は、ワーカーの障害に対する組み込みのフォールトトレランスを提供します。ワーカーが回復すると、以前に提供されたデータセット関数(カスタムトレーニングループの場合はClusterCoordinator.create_per_worker_dataset
、 Model.fit
の場合はtf.keras.utils.experimental.DatasetCreator
)がワーカーで呼び出され、データセットが再作成されます。
パラメータサーバーまたはコーディネータの障害
ただし、コーディネーターがパラメーターサーバーエラーを検出すると、すぐにUnavailableError
またはAbortedError
が発生します。この場合、コーディネーターを再起動できます。コーディネーター自体も利用できなくなる可能性があります。したがって、トレーニングの進行状況を失わないようにするために、特定のツールをお勧めします。
Model.fit
の場合、進行状況の保存と復元を自動的に処理するBackupAndRestore
コールバックを使用する必要があります。例については、上記のコールバックとトレーニングのセクションを参照してください。カスタムトレーニングループの場合、トレーニングを開始する前に、モデル変数を定期的にチェックポイントし、チェックポイントがある場合はチェックポイントからモデル変数をロードする必要があります。オプティマイザーがチェックポイントされている場合、トレーニングの進行状況は
optimizer.iterations
からおおよそ推測できます。
checkpoint_manager = tf.train.CheckpointManager(
tf.train.Checkpoint(model=model, optimizer=optimizer),
checkpoint_dir,
max_to_keep=3)
if checkpoint_manager.latest_checkpoint:
checkpoint = checkpoint_manager.checkpoint
checkpoint.restore(
checkpoint_manager.latest_checkpoint).assert_existing_objects_matched()
global_steps = int(optimizer.iterations.numpy())
starting_epoch = global_steps // steps_per_epoch
for _ in range(starting_epoch, num_epoches):
for _ in range(steps_per_epoch):
coordinator.schedule(step_fn, args=(per_worker_iterator,))
coordinator.join()
checkpoint_manager.save()
RemoteValue
の取得
関数が正常に実行された場合、 RemoteValue
のフェッチは成功することが保証されています。これは、現在、関数の実行後すぐに戻り値がコーディネーターにコピーされるためです。コピー中にワーカー障害が発生した場合、関数は別の使用可能なワーカーで再試行されます。したがって、パフォーマンスを最適化する場合は、戻り値なしで関数をスケジュールできます。
エラー報告
コーディネーターは、パラメーターサーバーからのUnavailableError
などのエラー、またはtf.debugging.check_numerics
からのInvalidArgument
などの他のアプリケーションエラーを検出すると、エラーを発生させる前に、保留中およびキューに入れられているすべての関数をキャンセルします。対応するRemoteValue
をフェッチすると、 CancelledError
が発生します。
エラーが発生した後、コーディネーターは同じエラーまたはキャンセルされた関数からのエラーを発生させません。
パフォーマンスの向上
ParameterServerStrategy
とClusterResolver
を使用してトレーニングするときにパフォーマンスの問題が発生する場合は、いくつかの理由が考えられます。
一般的な理由の1つは、パラメーターサーバーの負荷が不均衡であり、負荷の高いパラメーターサーバーの一部が容量に達していることです。複数の根本原因も考えられます。この問題を軽減するためのいくつかの簡単な方法は次のとおりです。
-
ParameterServerStrategy
を構築するときにvariable_partitioner
を指定して、大きなモデル変数をシャーディングします。 - 可能であれば、すべてのパラメータサーバーに必要なホットスポット変数を1つのステップで作成することは避けてください。たとえば、オプティマイザーで一定の学習率またはサブクラス
tf.keras.optimizers.schedules.LearningRateSchedule
を使用します。これは、デフォルトの動作では、学習率が特定のパラメーターサーバーに配置され、各ステップで他のすべてのパラメーターサーバーによって要求される変数になるためです。 。 - 大きな語彙をシャッフルしてから、Keras前処理レイヤーに渡します。
パフォーマンスの問題のもう1つの考えられる理由は、コーディネーターです。 schedule
/ join
の最初の実装はPythonベースであるため、スレッドのオーバーヘッドが発生する可能性があります。また、コーディネーターとワーカーの間のレイテンシーが大きくなる可能性があります。このような場合は、
Model.fit
の場合、Model.compile
で提供されるsteps_per_execution
引数を1より大きい値に設定できます。カスタムトレーニングループの場合、複数のステップを1つの
tf.function
にパックできます。
steps_per_invocation = 10
@tf.function
def step_fn(iterator):
for _ in range(steps_per_invocation):
features, labels = next(iterator)
def replica_fn(features, labels):
...
strategy.run(replica_fn, args=(features, labels))
ライブラリがさらに最適化されるため、将来、ほとんどのユーザーが手動でステップをパックする必要がなくなることを願っています。
さらに、パフォーマンスを向上させるための小さなトリックは、上記のタスク障害の処理のセクションで説明したように、戻り値なしで関数をスケジュールすることです。
既知の制限
既知の制限のほとんどは、上記のセクションですでに説明されています。このセクションでは、要約を示します。
ParameterServerStrategy
一般
- フォールトトレランスを適切に機能させるには、コーディネーターを含むすべてのタスクで
os.environment["grpc_fail_fast"]="use_caller"
が必要です。 - 同期パラメータサーバートレーニングはサポートされていません。
- 通常、最適なパフォーマンスを実現するには、複数のステップを1つの関数にまとめる必要があります。
- シャーディングされた変数を含むtf.saved_model.loadを介して
tf.saved_model.load
をロードすることはサポートされていません。 TensorFlowServingを使用してそのようなsaved_modelをロードすることは機能することが期待されていることに注意してください。 - シャーディングされたオプティマイザースロット変数を含むチェックポイントを異なる数のシャードにロードすることはサポートされていません。
- コーディネータータスクを再起動せずにパラメーターサーバーの障害から回復することはサポートされていません。
-
tf.lookup.StaticHashTable
(tf.keras.layers.IntegerLookup
、tf.keras.layers.StringLookup
、tf.keras.layers.TextVectorization
などの一部のKeras前処理レイヤーで一般的に使用されている)を使用すると、リソースが配置されますこの時点でのコーディネーターは、パラメーターサーバーのトレーニングを受けています。これは、ワーカーからコーディネーターへのルックアップRPCのパフォーマンスに影響します。これは、現在対処すべき最優先事項です。
Model.fit
の詳細
Model.fit
steps_per_epoch
引数が必要です。エポックで適切な間隔を提供する値を選択できます。-
ParameterServerStrategy
は、パフォーマンス上の理由から、バッチレベルの呼び出しを持つカスタムコールバックをサポートしていません。これらの呼び出しは、steps_per_epoch
のステップ数ごとに呼び出されるように、適切に選択されたsteps_per_epoch
を使用してエポックレベルの呼び出しに変換する必要があります。組み込みのコールバックは影響を受けません。バッチレベルの呼び出しは、パフォーマンスが向上するように変更されています。ParameterServerStrategy
のバッチレベルの呼び出しをサポートすることが計画されています。 - 同じ理由で、他の戦略とは異なり、プログレスバーとメトリックはエポック境界でのみログに記録されます。
-
run_eagerly
はサポートされていません。
カスタムトレーニングループの詳細
ClusterCoordinator.schedule
は、データセットの訪問保証をサポートしていません。-
ClusterCoordinator.create_per_worker_dataset
を使用する場合、データセット全体を、渡された関数内に作成する必要があります。 -
tf.data.Options
は、ClusterCoordinator.create_per_worker_dataset
によって作成されたデータセットでは無視されます。