You need to enable JavaScript to run this app.
导航

发起 TensorFlowPS 分布式训练

最近更新时间2021.12.14 18:04:53

首次发布时间2021.09.23 11:57:28

TensorFlow 支持在多台机器上进行参数服务器(parameter server)训练。在这种 ps-worker 的架构之下,部分实例会被指定为工作进程(worker),部分实例被指定为参数服务器(ps)。在每一轮训练中由 ps 将最新的模型参数分发给 worker,worker 完成计算后将梯度回传给 ps,ps 更新参数后再进入下一轮训练,直至训练完成。

基本流程
  1. 用户在【自定义训练】模块创建一个训练任务时选择实例配置为 TensorFlow PS,按需配置各种训练角色并提交任务表单进入任务创建环节。有如下几种训练角色:
    • ps:管理参数的存储和更新的参数服务器。实例数量必须 >= 1。
    • worker:执行训练任务的工作进程。在没有 chief 时,worker 的实例数量必须 >= 1。
    • chief:相当于 worker0,除了执行训练任务之外还负责保存 checkpoint 或日志。在没有 worker 时,chief 的实例数量必须 >= 1。
    • evaluator:执行评估任务。按需配置该实例的数量即可。
  2. 任务提交后,机器学习平台将为用户创建对应的实例并向所有实例注入相关的环境变量,通过环境变量用户代码得知集群的信息以及当前实例对应的训练角色,从而完成对应角色的本职任务直到训练结束。
    • 训练之前,平台将等待所有实例之间的网络通畅。
    • 训练过程中,任意实例失败(退出码非 0)则训练任务失败。
    • 所有实例训练完成(退出码为 0)则训练任务成功。
TF_CONFIG

TF_CONFIG 是 TensorFlow 分布式训练中最关键的环境变量,平台按照用户的实例配置并根据开源社区的规范向 TF_CONFIG 填充集群信息。如下是一个具体示例:

{
    "cluster": {
        "chief": [              # chief 的实例列表
            "chief0-host:port",
            "chief1-host:port"
            ...
        ],
        "evaluator": [         # evaluator 的实例列表
            "evaluator0-host:port",
            "evaluator1-host:port",
            ...
        ],
        "ps": [                # ps 的实例列表
            "ps0-host:port",
            "ps1-host:port",
            ...
        ],
        "worker": [            # worker 的实例列表
            "worker0-host:port",
            "worker1-host:port",
            ...
        ]
    },
    "environment": "cloud",
    "task": {
        "index": 0,           # 角色编号
        "type": "chief"       # 所属角色 
    }
}

每个实例根据 TF_CONFIG 中的 cluster 字段得知整个分布式集中 chief / worker / ps / evaluator 的数量以及对应的地址,同时从 task 字段中得知当前实例所属的角色以及编号。在该示例中当前实例的角色是 chief,在作为工作进程之外还需要完成 checkpoint 及日志的保存。

如需使用【自定义训练】模块的通用环境变量详见通用环境变量列表

示例代码

具体的示例代码详见 GitHub