Kubeflow实战系列: 利用TFJob运行分布式TensorFlow

介绍

本系列将介绍如何在阿里云容器服务上运行Kubeflow, 本文介绍如何使用TfJob运行分布式模型训练。

TensorFlow分布式训练和Kubernetes

TensorFlow作为现在最为流行的深度学习代码库,在数据科学家中间非常流行,特别是可以明显加速训练效率的分布式训练更是杀手级的特性。但是如何真正部署和运行大规模的分布式模型训练,却成了新的挑战。 实际分布式TensorFLow的使用者需要关心3件事情。

  1. 寻找足够运行训练的资源,通常一个分布式训练需要若干数量的worker(运算服务器)和ps(参数服务器),而这些运算成员都需要使用计算资源。
  2. 安装和配置支撑程序运算的软件和应用
  3. 根据分布式TensorFlow的设计,需要配置ClusterSpec。这个json格式的ClusterSpec是用来描述整个分布式训练集群的架构,比如需要使用两个worker和ps,ClusterSpec应该长成下面的样子,并且分布式训练中每个成员都需要利用这个ClusterSpec初始化tf.train.ClusterSpec对象,建立集群内部通信
cluster = tf.train.ClusterSpec({"worker": ["<VM_1>:2222",
                                           "<VM_2>:2222"],
                                "ps": ["<IP_VM_1>:2223",
                                       "<IP_VM_2>:2223"]})

其中第一件事情是Kubernetes资源调度非常擅长的事情,无论CPU和GPU调度,都是直接可以使用;而第二件事情是Docker擅长的,固化和可重复的操作保存到容器镜像。而自动化的构建ClusterSpecTFJob解决的问题,让用户通过简单的集中式配置,完成TensorFlow分布式集群拓扑的构建。

应该说烦恼了数据科学家很久的分布式训练问题,通过Kubernetes+TFJob的方案可以得到比较好的解决。

利用Kubernetes和TFJob部署分布式训练

  1. 修改TensorFlow分布式训练代码

之前在阿里云上小试TFJob一文中已经介绍了TFJob的定义,这里就不再赘述了。可以知道TFJob里有的角色类型为MASTERWORKER 和 PS

举个现实的例子,假设从事分布式训练的TFJob叫做distributed-mnist, 其中节点有1个MASTER, 2个WORKERS和2个PS,ClusterSpec对应的格式如下所示:

{  
    "master":[  
        "distributed-mnist-master-0:2222"
    ],
    "ps":[  
        "distributed-mnist-ps-0:2222",
        "distributed-mnist-ps-1:2222"
    ],
    "worker":[  
        "distributed-mnist-worker-0:2222",
        "distributed-mnist-worker-1:2222"
    ]
}

tf_operator的工作就是创建对应的5个Pod, 并且将环境变量TF_CONFIG传入到每个Pod中,TF_CONFIG包含三部分的内容,当前集群ClusterSpec, 该节点的角色类型,以及id。比如该Pod为worker0,它所收到的环境变量TF_CONFIG为:

{  
   "cluster":{  
      "master":[  
         "distributed-mnist-master-0:2222"
      ],
      "ps":[  
         "distributed-mnist-ps-0:2222"
      ],
      "worker":[  
         "distributed-mnist-worker-0:2222",
         "distributed-mnist-worker-1:2222"
      ]
   },
   "task":{  
      "type":"worker",
      "index":0
   },
   "environment":"cloud"
}

在这里,tf_operator负责将集群拓扑的发现和配置工作完成,免除了使用者的麻烦。对于使用者来说,他只需要在这里代码中使用通过获取环境变量TF_CONFIG中的上下文。

这意味着,用户需要根据和TFJob的规约修改分布式训练代码:

# 从环境变量TF_CONFIG中读取json格式的数据
tf_config_json = os.environ.get("TF_CONFIG", "{}")

# 反序列化成python对象
tf_config = json.loads(tf_config_json)

# 获取Cluster Spec
cluster_spec = tf_config.get("cluster", {})
cluster_spec_object = tf.train.ClusterSpec(cluster_spec)

# 获取角色类型和id, 比如这里的job_name 是 "worker" and task_id 是 0
task = tf_config.get("task", {})
job_name = task["type"]
task_id = task["index"]

# 创建TensorFlow Training Server对象
server_def = tf.train.ServerDef(
    cluster=cluster_spec_object.as_cluster_def(),
    protocol="grpc",
    job_name=job_name,
    task_index=task_id)
server = tf.train.Server(server_def)

# 如果job_name为ps,则调用server.join()
if job_name == 'ps':
    server.join()

# 检查当前进程是否是master, 如果是master,就需要负责创建session和保存summary。
is_chief = (job_name == 'master')


# 通常分布式训练的例子只有ps和worker两个角色,而在TFJob里增加了master这个角色,实际在分布式TensorFlow的编程模型并没有这个设计。而这需要使用TFJob的分布式代码里进行处理,不过这个处理并不复杂,只需要将master也看做worker_device的类型
with tf.device(tf.train.replica_device_setter(
    worker_device="/job:{0}/task:{1}".format(job_name,task_id),
    cluster=cluster_spec)):

具体代码可以参考示例代码

2. 在本例子中,将演示如何使用TFJob运行分布式训练,并且将训练结果和日志保存到NAS存储上,最后通过Tensorboard读取训练日志。

2.1 创建NAS数据卷,并且设置与当前Kubernetes集群的同一个具体vpc的挂载点。操作详见文档

2.2 在NAS上创建 /training的数据文件夹, 下载mnist训练所需要的数据

原文链接

阅读更多
想对作者说点什么?

博主推荐

换一批

没有更多推荐了,返回首页