Google I/O にご注目いただきありがとうございます。すべてのセッションをオンデマンドで表示オンデマンドで見る

ParameterServerStrategyを使用したパラメーターサーバーのトレーニング

TensorFlow.orgで表示GoogleColabで実行GitHubでソースを表示 ノートブックをダウンロード

概要

パラメータサーバートレーニングは、複数のマシンでモデルトレーニングをスケールアップするための一般的なデータ並列方法です。

パラメーターサーバートレーニングクラスターは、ワーカーとパラメーターサーバーで構成されます。変数はパラメーターサーバー上に作成され、各ステップでワーカーによって読み取られて更新されます。デフォルトでは、ワーカーはこれらの変数を相互に同期せずに個別に読み取り、更新します。これが、パラメータサーバースタイルのトレーニングが非同期トレーニングと呼ばれることがある理由です。

TensorFlow 2では、パラメーターサーバーのトレーニングはtf.distribute.experimental.ParameterServerStrategyクラスによって強化されます。このクラスは、トレーニングステップを数千人のワーカー（パラメーターサーバーを伴う）にスケールアップするクラスターに分散します。

サポートされているトレーニング方法

サポートされている主なトレーニング方法は2つあります。

ジョブとタスクを含むクラスター

選択した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

クラスターのセットアップ

上記のように、パラメーターサーバートレーニングクラスターには、トレーニングプログラムを実行するコーディネータータスク、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つは、パラメーターサーバーの負荷が不均衡であり、負荷の高いパラメーターサーバーの一部が容量に達していることです。複数の根本原因も考えられます。この問題を軽減するためのいくつかの簡単な方法は次のとおりです。

  1. ParameterServerStrategyを構築するときにvariable_partitionerを指定して、大きなモデル変数をシャーディングします。
  2. 可能であれば、すべてのパラメータサーバーに必要なホットスポット変数を1つのステップで作成することは避けてください。たとえば、オプティマイザーで一定の学習率またはサブクラスtf.keras.optimizers.schedules.LearningRateScheduleを使用します。これは、デフォルトの動作では、学習率が特定のパラメーターサーバーに配置され、各ステップで他のすべてのパラメーターサーバーによって要求される変数になるためです。 。
  3. 大きな語彙をシャッフルしてから、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はサポートされていません。

カスタムトレーニングループの詳細