TF_CONFIG 和分布式训练

运行训练作业时，AI Platform Training 会在作业中的每个虚拟机实例上设置一个名为 TF_CONFIG 的环境变量。在每个虚拟机上运行的训练代码可以使用 TF_CONFIG 环境变量访问训练作业的详细信息及其运行的虚拟机角色。

TensorFlow 使用 TF_CONFIG 环境变量促进分布式训练，但您可能无需直接在训练代码中访问它。本文档介绍了 TF_CONFIG 环境变量及其在分布式 TensorFlow 作业和超参数调节作业中的用法。

TF_CONFIG 的格式

AI Platform Training 在每个训练作业的每个虚拟机上设置 TF_CONFIG 环境变量，以满足 TensorFlow 对分布式训练要求的规范。然而，AI Platform Training 还在 TF_CONFIG 环境变量中设置超出 TensorFlow 要求的其他字段。

TF_CONFIG 环境变量是一个 JSON 字符串，格式如下：

TF_CONFIG 字段
cluster

TensorFlow 集群描述。将一个或多个任务名称（chief、worker、ps 或 master）映射到运行这些任务的网络地址列表的字典。对于给定的训练作业，此字典在每个虚拟机上都是相同的。

这是 tf.train.ClusterSpec 构造函数的有效第一个参数。请注意，此字典不会包含 evaluator 作为键，因为即使您将评估程序用于作业，评估程序也不会被视为训练集群的一部分。

请在本文档的其他部分中了解 chief 和 master 之间的区别。
task

设置此环境变量的虚拟机的任务说明。 对于给定的训练作业，此字典在每个虚拟机上都是不同的。您可以使用此信息自定义分布式训练作业中每个虚拟机上运行的代码。您还可以使用它来更改训练代码的行为，使其适应超参数调节作业的不同试验。

此字典包含以下键值对：

task 字段
type

此虚拟机正在执行的任务类型。此值在工作器中设置为 worker，在参数服务器上设置为 ps，而在评估程序上设置为 evaluator。在作业的主工作器中，值设置为 chief 或 master；可在本文档的 chief 对比 master 部分详细了解两者之间的区别。
index

任务的索引（从零开始）。例如，如果训练作业包含两个工作器，则其中一个值设为 0，另一个设为 1。
trial

当前在此虚拟机上运行的超参数调节试验的 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 或更高版本中运行的超参数调节作业使用一个主工作器、两个工作器和一个参数服务器，此代码在该作业中在第一次超参数调节试验期间为其中一个工作器生成以下日志。示例输出为了简洁会隐藏 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 中的主工作器虚拟机与 TensorFlow 中的 chief 任务类型相对应。虽然 TensorFlow 可以指定 worker 任务作为 chief，但 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 将使用已弃用的 master 任务类型取代 chief。

何时使用 TF_CONFIG

如上一节所述，您可能不需要直接在训练代码中与 TF_CONFIG 环境变量互动。仅在 TensorFlow 的分布策略和 AI Platform Training 的标准超参数调节工作流程（下节会对二者进行介绍）都不适用于您的作业的情况下，才能访问 TF_CONFIG 环境变量。

分布式训练

AI Platform Training 设置 TF_CONFIG 环境变量来扩展 TensorFlow 分布式训练所需的规范。

如要使用 TensorFlow 执行分布式训练，请使用 tf.distribute.Strategy API。我们尤其建议您将 Keras API 与 MultiWorkerMirroredStrategy 结合使用，或者，如果您为作业指定参数服务器，那么将 Keras API 与 ParameterServerStrategy 结合使用。不过请注意，TensorFlow 目前仅为这些策略提供实验性支持。

这些分布策略使用 TF_CONFIG 环境变量为训练作业中的每个虚拟机分配角色，并促进虚拟机之间的通信。您无需直接在训练代码中访问 TF_CONFIG 环境变量，因为 TensorFlow 会为您处理。

如果要自定义运行训练作业的不同虚拟机的行为方式，仅需直接解析 TF_CONFIG 环境变量。

超参数调节

当您运行超参数调节作业时，AI Platform Training 会为每次试验的训练代码提供不同的参数。您的训练代码不必知道当前正在运行的试验。此外，AI Platform Training 提供了监控超参数调节作业进度的工具。

如果需要，您的代码可以从 TF_CONFIG 环境变量的 task 字段的 trial 字段中读取当前的试验编号。

后续步骤