本页面介绍了自定义容器如何支持 AI Platform Training 上的分布式训练结构。
借助自定义容器,您可以使用任何支持分布的机器学习框架进行分布式训练。虽然此处使用的术语是以 TensorFlow 的分布式模型为基础,但您可以使用其他任何具有类似分布式结构的机器学习框架。例如,MXNet 中的分布式训练使用调度器、工作器和服务器。这与 AI Platform Training 自定义容器的分布式训练结构相对应;后者使用主实例、工作器和参数服务器。
训练集群的结构
如果您使用 AI Platform Training 运行分布式训练作业,则需要在训练集群中指定多个机器(节点)。训练服务会为您指定的机器类型分配资源。给定节点上正在运行的作业称为副本。根据分布式 TensorFlow 模型,训练集群中的每个副本将在分布式训练中获得单个角色或任务:
主实例工作器:您只能将一个副本指定为主实例工作器(也称为主工作器)。这项任务可管理其他任务,并报告作业的整体状态。
工作器:您可以将一个或多个副本指定为工作器。这些副本执行您在作业配置中为它们指定的工作。
参数服务器:可以将一个或多个副本指定为参数服务器。这些副本会存储模型参数,并协调工作器之间的共享模型状态。
评估程序:您可以将一个或多个副本指定为评估器。这些副本可用于评估您的模型。如果您使用的是 TensorFlow,请注意 TensorFlow 通常要求您只能使用一个评估器。
API 映射
您可以为训练集群中的机器分配四种不同角色,这些角色对应于您可以在 TrainingInput
(表示训练作业的输入参数)中指定的四个字段:
masterConfig.imageUri
表示将要在主实例工作器上运行的容器映像 URI。workerConfig.imageUri
、parameterServerConfig.imageUri
和evaluatorConfig.imageUri
分别表示将要在工作器、参数服务器和评估器上运行的容器映像 URI。如果没有为这些字段设置值,AI Platform Training 会使用masterConfig.imageUri
的值。
您还可以在 gcloud ai-platform jobs submit training
中为每个字段设置相应标志的值:
- 对于主实例工作器配置,请使用
--master-image-uri
。 - 对于工作器配置,请使用
--worker-image-uri
。 - 对于参数服务器配置,请使用
--parameter-server-image-uri
。 - 目前没有用于为评估器指定容器映像 URI 的标志。您可以在 config.yaml 配置文件中指定
evaluatorConfig.imageUri
。
请参阅如何使用自定义容器提交分布式训练作业的示例。
了解 CLUSTER_SPEC
AI Platform Training 会针对每个副本填充一个 CLUSTER_SPEC
环境变量,以便描述整个集群的设置方式。与 TensorFlow 的 TF_CONFIG
类似,CLUSTER_SPEC
描述集群中的每个副本,包括其索引和角色(主实例工作器、工作器、参数服务器或评估器)。
当您使用 TensorFlow 运行分布式训练时,系统会解析 TF_CONFIG
以构建 tf.train.ClusterSpec
。同样,当您使用其他机器学习框架运行分布式训练时,必须解析 CLUSTER_SPEC
以填充框架所需的任何环境变量或设置。
CLUSTER_SPEC
的格式
CLUSTER_SPEC
环境变量是一个 JSON 字符串,格式如下:
键 | 说明 | |
---|---|---|
"cluster"
|
自定义容器的集群说明。与“TF_CONFIG”一样,此对象也会采用 TensorFlow 集群规范中的格式,且可以传递给 tf.train.ClusterSpec 的构造函数。 | |
"task"
|
描述代码运行所在特定节点的任务。您可以使用此信息为分布式作业中的特定工作器编写代码。此条目是一个含有以下键的字典: | |
"type"
|
此节点执行的任务类型。可能的值包括 master 、worker 、ps 和 evaluator 。 |
|
"index"
|
任务的索引(从零开始)。大多数分布式训练作业都有一个主实例任务、一个或多个参数服务器,以及一个或多个工作器。 | |
"trial"
|
当前运行的超参数调节试验的标识符。为作业配置超参数调节时,可以设置多个要训练的试验。此值为您提供了一种区分代码中正在运行的试验的方法。标识符是一个包含试验编号的字符串值,从 1 开始。 | |
"job"
|
启动作业时使用的作业参数。在大多数情况下,您可以忽略此条目,因为其中包含的数据与通过命令行参数传递给应用的数据完全一致。 |
与 TF_CONFIG
比较
请注意,AI Platform Training 还会对所有训练作业的每个副本设置 TF_CONFIG
环境变量。AI Platform Training 仅会对自定义容器训练作业的副本设置 CLUSTER_SPEC
环境变量。这两个环境变量共用一些值,但格式不同。
使用自定义容器训练时,默认情况下主实例副本在 TF_CONFIG
环境变量中标有 master
任务名称。您可以将主实例副本配置为标有 chief
任务名称,方法是在创建训练作业时将 trainingInput.useChiefInTfConfig
字段设置为 true
,或者在作业中使用一个或多个评估器副本。这在您的自定义容器使用 TensorFlow 2 的情况下尤为有用。
除了此配置选项之外,在您使用自定义容器时与在您使用 AI Platform Training 运行时版本时,使用 TensorFlow 进行分布式训练的方法是相同的。查看有关如何在 AI Platform Training 上使用 TF_CONFIG
进行分布式训练的详细信息和示例。