利用TFJob运行分布式TensorFlow
TensorFlow分布式训练和Kubernetes
TensorFlow作为现在最为流行的深度学习代码库,在数据科学家中间非常流行,特别是可以明显加速训练效率的分布式训练更是杀手级的特性。但是如何真正部署和运行大规模的分布式模型训练,却成了新的挑战。 实际分布式TensorFLow的使用者需要关心3件事情。
- 寻找足够运行训练的资源,通常一个分布式训练需要若干数量的worker(运算服务器)和ps(参数服务器),而这些运算成员都需要使用计算资源。
- 安装和配置支撑程序运算的软件和应用
- 根据分布式TensorFlow的设计,需要配置ClusterSpec。这个json格式的ClusterSpec是用来描述整个分布式训练集群的架构,比如需要使用两个worker和ps,ClusterSpec应该长成下面的样子,并且分布式训练中每个成员都需要利用这个ClusterSpec初始化tf.train.ClusterSpec对象,建立集群内部通信
其中第一件事情是Kubernetes资源调度非常擅长的事情,无论CPU和GPU调度,都是直接可以使用;而第二件事情是Docker擅长的,固化和可重复的操作保存到容器镜像。而自动化的构建ClusterSpec是TFJob解决的问题,让用户通过简单的集中式配置,完成TensorFlow分布式集群拓扑的构建。
应该说烦恼了数据科学家很久的分布式训练问题,通过Kubernetes+TFJob的方案可以得到比较好的解决。
利用Kubernetes和TFJob部署分布式训练
1.修改TensorFlow分布式训练代码
之前在阿里云上小试TFJob一文中已经介绍了TFJob的定义,这里就不再赘述了。可以知道TFJob里有的角色类型为MASTER, WORKER 和 PS。
举个现实的例子,假设从事分布式训练的TFJob叫做distributed-mnist, 其中节点有1个MASTER, 2个WORKERS和2个PS,ClusterSpec对应的格式如下所示:
而tf_operator的工作就是创建对应的5个Pod, 并且将环境变量TF_CONFIG传入到每个Pod中,TF_CONFIG包含三部分的内容,当前集群ClusterSpec, 该节点的角色类型,以及id。比如该Pod为worker0,它所收到的环境变量TF_CONFIG为:
在这里,tf_operator负责将集群拓扑的发现和配置工作完成,免除了使用者的麻烦。对于使用者来说,他只需要在这里代码中使用通过获取环境变量TF_CONFIG中的上下文。
这意味着,用户需要根据和TFJob的规约修改分布式训练代码:
具体代码可以参考示例代码
2. 在本例子中,将演示如何使用TFJob运行分布式训练,并且将训练结果和日志保存到NAS存储上,最后通过Tensorboard读取训练日志。
2.1 创建NAS数据卷,并且设置与当前Kubernetes集群的同一个具体vpc的挂载点。操作详见文档
2.2 在NAS上创建 /training的数据文件夹, 下载mnist训练所需要的数据
2.3 创建NAS的PV, 以下为示例nas-dist-pv.yaml
将该模板保存到nas-dist-pv.yaml, 并且创建pv:
2.4 利用nas-dist-pvc.yaml创建PVC
具体命令:
2.5 创建TFJob
将该模板保存到mnist-simple-gpu-dist.yaml, 并且创建分布式训练的TFJob:
检查所有运行的Pod
查看master的日志,可以看到ClusterSpec已经成功的构建出来了
2.6 部署TensorBoard,并且查看训练效果
为了更方便 TensorFlow 程序的理解、调试与优化,可以用 TensorBoard 来观察 TensorFlow 训练效果,理解训练框架和优化算法, 而TensorBoard通过读取TensorFlow的事件日志获取运行时的信息。
在之前的分布式训练样例中已经记录了事件日志,并且保存到文件events.out.tfevents*中
在Kubernetes部署TensorBoard, 并且指定之前训练的NAS存储
将该模板保存到tensorboard.yaml, 并且创建tensorboard:
TensorBoard创建成功后,通过kubectl port-forward命令进行访问
通过http://127.0.0.1:6006登录TensorBoard,查看分布式训练的模型和效果:
总结
利用tf-operator可以解决分布式训练的问题,简化数据科学家进行分布式训练工作。同时使用Tensorboard查看训练效果, 再利用NAS或者OSS来存放数据和模型,这样一方面有效的重用训练数据和保存实验结果,另外一方面也是为模型预测的发布做准备。如何把模型训练,验证,预测串联起来构成机器学习的工作流(workflow), 也是Kubeflow的核心价值,我们在后面的文章中也会进行介绍。