运行训练作业时,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 集群描述。将一个或多个任务名称( 这是 请在本文档的其他部分中了解 |
||||||||||
task |
设置此环境变量的虚拟机的任务说明。 对于给定的训练作业,此字典在每个虚拟机上都是不同的。您可以使用此信息自定义分布式训练作业中每个虚拟机上运行的代码。您还可以使用它来更改训练代码的行为,使其适应超参数调节作业的不同试验。 此字典包含以下键值对:
|
||||||||||
job |
您创建当前训练作业时提供的 |
||||||||||
environment |
字符串 |
对于自定义容器训练作业,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 中的主工作器虚拟机对应于
到chief
任务
类型
TensorFlow 中的代码。虽然 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
:
- 您正在运行使用运行时版本 2.1 或更高版本的训练作业。
- 您已将您的训练作业配置为使用一个或多个评估程序。换句话说,您已将作业的
trainingInput.evaluatorCount
设置为1
或更大。 - 您的作业使用自定义容器,并且您已将您的作业的
trainingInput.useChiefInTfConfig
设置为true
。
否则,出于兼容性考虑,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
字段中读取当前的试验编号。
后续步骤
- 通过 TensorFlow 文档中关于使用 Keras 的多工作器训练教程进行工作。
- 了解如何在 AI Platform Training 中使用自定义容器进行分布式训练。
- 了解如何为训练作业实现超参数调节。