TF_CONFIG と分散トレーニング

トレーニング ジョブを実行すると、AI Platform Training によって、ジョブの一部である各仮想マシン（VM）インスタンスに TF_CONFIG という環境変数が設定されます。各 VM で実行されるトレーニング コードは、TF_CONFIG 環境変数を使用して、実行中の VM のトレーニング ジョブとロールの詳細にアクセスできます。

TensorFlow は、分散トレーニングを容易にするために TF_CONFIG 環境変数を使用しますが、トレーニング コードで直接アクセスする必要はありません。このドキュメントでは、分散 TensorFlow ジョブおよびハイパーパラメータ調整ジョブにおける TF_CONFIG 環境変数とその使用方法について説明します。

TF_CONFIG の形式

AI Platform Training では、TensorFlow が分散トレーニングに要求する仕様に合わせて、すべてのトレーニング ジョブのすべての VM で TF_CONFIG 環境変数を設定します。ただし、AI Platform Training では、TF_CONFIG 環境変数に TensorFlow で必要とするもの以外のフィールドも追加して設定します。

TF_CONFIG 環境変数は、次の形式の JSON 文字列です。

TF_CONFIG フィールド
cluster

TensorFlow クラスタの説明。1 つ以上のタスク名（chief、worker、ps、または master）を、タスクが実行されているネットワーク アドレスのリストにマッピングする辞書。特定のトレーニング ジョブでは、この辞書はすべての VM で同じです。

これは tf.train.ClusterSpec コンストラクタの有効な 1 番目の引数です。評価ツールをジョブで使用していても、評価ツールがトレーニング クラスタの一部と見なされることはないため、この辞書に evaluator がキーとして含まれることはありません。

chief と master の違いについては、このドキュメントの別のセクションをご覧ください。
task

この環境変数が設定されている VM のタスクの説明。特定のトレーニング ジョブでは、この辞書は VM によって違います。この情報を使用して、分散トレーニング ジョブの各 VM で実行するコードをカスタマイズできます。また、この情報を使用して、ハイパーパラメータ調整ジョブのさまざまなトライアルを実行する際に、トレーニング コードの動作を変更することもできます。

この辞書には、次の Key-Value ペアが含まれています。

task フィールド
type

この VM で実行されているタスクのタイプ。この値はワーカーでは worker、パラメータ サーバーでは ps、評価ツールでは evaluator に設定されます。ジョブのマスター ワーカーでは、値は chief または master のいずれかに設定されます。この 2 つの違いについては、このドキュメントの chief 対 master のセクションをご覧ください。
index

ゼロから開始するタスクのインデックス。たとえば、トレーニング ジョブに 2 つのワーカーが含まれている場合、この値は 1 つのワーカーの 0 と、もう 1 つのワーカーの 1 に設定されます。
trial

この VM で現在実行中のハイパーパラメータ調整トライアルの ID。このフィールドは、現在のトレーニング ジョブがハイパーパラメータ調整ジョブである場合にのみ設定されます。

ハイパーパラメータ調整ジョブの場合、AI Platform Training によって毎回さまざまなハイパーパラメータで多数のトライアルが行われ、トレーニング コードが繰り返し実行されます。このフィールドには現在のトライアル番号が含まれます。トライアル番号は最初のトライアルで 1 から始まります。
cloud

AI Platform Training で内部的に使用される ID。このフィールドは無視できます。
job

現在のトレーニング ジョブを作成するために提供した TrainingInput で、辞書として表現されます。
environment

文字 cloud。

カスタム コンテナのトレーニング ジョブの場合、AI Platform Training では CLUSTER_SPEC という環境変数を追加で設定します。この環境変数の形式は TF_CONFIG に似ていますが、重要な違いがいくつかあります。詳しくは、CLUSTER_SPEC 環境変数をご覧ください。

例

次のサンプルコードでは、TF_CONFIG 環境変数をトレーニング ログに出力します。

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

ハイパーパラメータ調整ジョブがランタイム バージョン 2.1 以降で、マスター ワーカー、2 つのワーカー、パラメータ サーバーを使用している場合、このコードでは最初のハイパーパラメータ調整トライアル中に、ワーカーのいずれかで次のログを出力します。出力例では、簡潔にするために job フィールドを非表示にし、一部の ID を汎用値に置き換えています。

{
  "cluster": {
    "chief": [
      "cmle-training-chief-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "cmle-training-ps-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "cmle-training-worker-[ID_STRING_1]-0:2222",
      "cmle-training-worker-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

chief 対 master

AI Platform Training のマスター ワーカー VM は、TensorFlow の chief タスクタイプに対応しています。TensorFlow では chief として機能する worker タスクを指定できますが、AI Platform Training では常に chief が明示的に指定されます。

master は TensorFlow の非推奨のタスクタイプです。master は同様のロールを chief として実行したタスクを表していましたが、一部の構成では、evaluator としても機能しました。TensorFlow 2 では、master タスクを含む TF_CONFIG 環境変数をサポートしていません。

次のいずれかに該当する場合、AI Platform Training は TF_CONFIG 環境変数の cluster フィールドと task フィールドで chief を使用します。

それ以外の場合は互換性を保つために、AI Platform Training では chief の代わりに非推奨の master タスクタイプを使用します。

TF_CONFIG を使用するタイミング

前のセクションで説明したように、トレーニング コードで TF_CONFIG 環境変数を直接操作する必要はなさそうです。TF_CONFIG 環境変数にアクセスするのは、TensorFlow の分散戦略と AI Platform Training の標準のハイパーパラメータ調整ワークフロー（両方とも以降のセクションで説明します）がジョブで機能しない場合のみです。

分散トレーニング

AI Platform Training では、TensorFlow が分散トレーニングに必要とする仕様を拡張するように、TF_CONFIG 環境変数を設定します。

TensorFlow を使用して分散トレーニングを実行するには、tf.distribute.Strategy API を使用します。特に、Keras API を MultiWorkerMirroredStrategy または ParameterServerStrategy（ジョブのパラメータ サーバーを指定する場合）と使用することをおすすめします。ただし、TensorFlow は現在、これらの戦略に対して試験運用版のサポートのみを提供していることに注意してください。

これらの分散戦略では、環境変数 TF_CONFIG を使用して、トレーニング ジョブの各 VM にロールを割り当て、VM 間の通信を容易にします。TensorFlow が環境変数 TF_CONFIG を処理するため、トレーニング コードで直接アクセスする必要はありません。

環境変数 TF_CONFIG を直接解析するのは、トレーニング ジョブを実行する各 VM の動作をカスタマイズする場合のみです。

ハイパーパラメータ調整

ハイパーパラメータ調整ジョブを実行すると、AI Platform Training から各トライアルのトレーニング コードにさまざまな引数が渡されます。トレーニング コードでは、現在どのトライアルを実行中かということを必ずしも認識する必要はありません。さらに、AI Platform Training には、ハイパーパラメータ調整ジョブの進行状況をモニタリングするためのツールが用意されています。

必要に応じて、コードで TF_CONFIG 環境変数の task フィールドの trial フィールドから、現在のトライアル番号を読み取ることができます。

次のステップ