摘要: TensorFlow 作为现在最为流行的深度学习代码库, 在数据科学家中间非常流行, 特别是可以明显加速训练效率的分布式训练更是杀手级的特性. 但是如何真正部署和运行大规模的分布式模型训练, 却成了新的挑战.
介绍
本系列将介绍如何在阿里云容器服务上运行 Kubeflow, 本文介绍如何使用 TfJob 运行分布式模型训练.
TensorFlow 分布式训练和 Kubernetes
TensorFlow 作为现在最为流行的深度学习代码库, 在数据科学家中间非常流行, 特别是可以明显加速训练效率的分布式训练更是杀手级的特性. 但是如何真正部署和运行大规模的分布式模型训练, 却成了新的挑战. 实际分布式 TensorFLow 的使用者需要关心 3 件事情.
寻找足够运行训练的资源, 通常一个分布式训练需要若干数量的 worker(运算服务器) 和 ps(参数服务器), 而这些运算成员都需要使用计算资源.
安装和配置支撑程序运算的软件和应用
根据分布式 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 擅长的, 固化和可重复的操作保存到容器镜像. 而自动化的构建 ClusterSpec 是 TFJob 解决的问题, 让用户通过简单的集中式配置, 完成 TensorFlow 分布式集群拓扑的构建.
应该说烦恼了数据科学家很久的分布式训练问题, 通过 Kubernetes+TFJob 的方案可以得到比较好的解决.
利用 Kubernetes 和 TFJob 部署分布式训练
修改 TensorFlow 分布式训练代码
之前在阿里云上小试 TFJob 一文中已经介绍了 TFJob 的定义, 这里就不再赘述了. 可以知道 TFJob 里有的角色类型为 MASTER, WORKER 和 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)):
具体代码可以参考示例代码
在本例子中, 将演示如何使用 TFJob 运行分布式训练, 并且将训练结果和日志保存到 NAS 存储上, 最后通过 Tensorboard 读取训练日志.
2.1 创建 NAS 数据卷, 并且设置与当前 Kubernetes 集群的同一个具体 vpc 的挂载点. 操作详见文档
2.2 在 NAS 上创建 /training 的数据文件夹, 下载 mnist 训练所需要的数据
mkdir -p /nfs
mount -t nfs -o vers=4.0 xxxxxxx.cn-hangzhou.nas.aliyuncs.com://nfs
- mkdir -p /nfs/training
- umount /nfs
2.3 创建 NAS 的 PV, 以下为示例 nas-dist-pv.yaml
- apiVersion: v1
- kind: PersistentVolume
- metadata:
- name: kubeflow-dist-nas-mnist
- labels:
- tfjob: kubeflow-dist-nas-mnist
- spec:
- capacity:
- storage: 10Gi
- accessModes:
- ReadWriteMany
- storageClassName: nas
- flexVolume:
- driver: "alicloud/nas"
- options:
- mode: "755"
- path: /training
- server: xxxxxxx.cn-hangzhou.nas.aliyuncs.com
- vers: "4.0"
将该模板保存到 nas-dist-pv.yaml, 并且创建 pv:
- kubectl create -f nas-dist-pv.yaml
- persistentvolume "kubeflow-dist-nas-mnist" created
2.4 利用 nas-dist-pvc.yaml 创建 PVC
- kind: PersistentVolumeClaim
- apiVersion: v1
- metadata:
- name: kubeflow-dist-nas-mnist
- spec:
- storageClassName: nas
- accessModes:
- ReadWriteMany
- resources:
- requests:
- storage: 5Gi
- selector:
- matchLabels:
- tfjob: kubeflow-dist-nas-mnist
具体命令:
- kubectl create -f nas-dist-pvc.yaml
- persistentvolumeclaim "kubeflow-dist-nas-mnist" created
2.5 创建 TFJob
- apiVersion: kubeflow.org/v1alpha1
- kind: TFJob
- metadata:
- name: mnist-simple-gpu-dist
- spec:
- replicaSpecs:
- replicas: 1 # 1 Master
- tfReplicaType: MASTER
- template:
- spec:
- containers:
- image: registry.aliyuncs.com/tensorflow-samples/tf-mnist-distributed:gpu
- name: tensorflow
- env:
- name: TEST_TMPDIR
- value: /training
- command: ["python", "/app/main.py"]
- resources:
- limits:
- nvidia.com/gpu: 1
- volumeMounts:
- name: kubeflow-dist-nas-mnist
- mountPath: "/training"
- volumes:
- name: kubeflow-dist-nas-mnist
- persistentVolumeClaim:
- claimName: kubeflow-dist-nas-mnist
- restartPolicy: OnFailure
- replicas: 1 # 1 or 2 Workers depends on how many gpus you have
- tfReplicaType: WORKER
- template:
- spec:
- containers:
- image: registry.aliyuncs.com/tensorflow-samples/tf-mnist-distributed:gpu
- name: tensorflow
- env:
- name: TEST_TMPDIR
- value: /training
- command: ["python", "/app/main.py"]
- imagePullPolicy: Always
- resources:
- limits:
- nvidia.com/gpu: 1
- volumeMounts:
- name: kubeflow-dist-nas-mnist
- mountPath: "/training"
- volumes:
- name: kubeflow-dist-nas-mnist
- persistentVolumeClaim:
- claimName: kubeflow-dist-nas-mnist
- restartPolicy: OnFailure
- replicas: 1 # 1 Parameter server
- tfReplicaType: PS
- template:
- spec:
- containers:
- image: registry.aliyuncs.com/tensorflow-samples/tf-mnist-distributed:cpu
- name: tensorflow
- command: ["python", "/app/main.py"]
- env:
- name: TEST_TMPDIR
- value: /training
- imagePullPolicy: Always
- volumeMounts:
- name: kubeflow-dist-nas-mnist
- mountPath: "/training"
- volumes:
- name: kubeflow-dist-nas-mnist
- persistentVolumeClaim:
- claimName: kubeflow-dist-nas-mnist
- restartPolicy: OnFailure
将该模板保存到 mnist-simple-gpu-dist.yaml, 并且创建分布式训练的 TFJob:
kubectl create -f mnist-simple-gpu-dist.yaml
tfjob "mnist-simple-gpu-dist" created
检查所有运行的 Pod
- RUNTIMEID=$(kubectl get tfjob mnist-simple-gpu-dist -o=jsonpath='{.spec.RuntimeId}')
- kubectl get po -lruntime_id=$RUNTIMEID
NAME READY STATUS RESTARTS AGE
mnist-simple-gpu-dist-master-z5z4-0-ipy0s 1/1 Running 0 31s
mnist-simple-gpu-dist-ps-z5z4-0-3nzpa 1/1 Running 0 31s
mnist-simple-gpu-dist-worker-z5z4-0-zm0zm 1/1 Running 0 31s
查看 master 的日志, 可以看到 ClusterSpec 已经成功的构建出来了
- kubectl logs -l runtime_id=$RUNTIMEID,job_type=MASTER
- 2018-06-10 09:31:55.342689: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1105] Found device 0 with properties:
name: Tesla P100-PCIE-16GB major: 6 minor: 0 memoryClockRate(GHz): 1.3285
pciBusID: 0000:00:08.0
totalMemory: 15.89GiB freeMemory: 15.60GiB
2018-06-10 09:31:55.342724: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1195] Creating TensorFlow device (/device:GPU:0) -> (device: 0, name: Tesla P100-PCIE-16GB, pci bus id: 0000:00:08.0, compute capability: 6.0)
- 2018-06-10 09:31:55.805747: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job master -> {0 -> localhost:2222}
- 2018-06-10 09:31:55.805786: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job ps -> {0 -> mnist-simple-gpu-dist-ps-m5yi-0:2222}
- 2018-06-10 09:31:55.805794: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:215] Initialize GrpcChannelCache for job worker -> {0 -> mnist-simple-gpu-dist-worker-m5yi-0:2222}
- 2018-06-10 09:31:55.807119: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:324] Started server with target: grpc://localhost:2222
- ...
- Accuracy at step 900: 0.9709
- Accuracy at step 910: 0.971
- Accuracy at step 920: 0.9735
- Accuracy at step 930: 0.9716
- Accuracy at step 940: 0.972
- Accuracy at step 950: 0.9697
- Accuracy at step 960: 0.9718
- Accuracy at step 970: 0.9738
- Accuracy at step 980: 0.9725
- Accuracy at step 990: 0.9724
- Adding run metadata for 999
2.6 部署 TensorBoard, 并且查看训练效果
为了更方便 TensorFlow 程序的理解, 调试与优化, 可以用 TensorBoard 来观察 TensorFlow 训练效果, 理解训练框架和优化算法, 而 TensorBoard 通过读取 TensorFlow 的事件日志获取运行时的信息.
在之前的分布式训练样例中已经记录了事件日志, 并且保存到文件 events.out.tfevents * 中
- tree
- .
- tensorflow
- input_data
- t10k-images-idx3-ubyte.gz
- t10k-labels-idx1-ubyte.gz
- train-images-idx3-ubyte.gz
- train-labels-idx1-ubyte.gz
- logs
- checkpoint
events.out.tfevents.1528760350.mnist-simple-gpu-dist-master-fziz-0-74je9
- graph.pbtxt
- model.ckpt-0.data-00000-of-00001
- model.ckpt-0.index
- model.ckpt-0.meta
- test
events.out.tfevents.1528760351.mnist-simple-gpu-dist-master-fziz-0-74je9
events.out.tfevents.1528760356.mnist-simple-gpu-dist-worker-fziz-0-9mvsd
train
events.out.tfevents.1528760350.mnist-simple-gpu-dist-master-fziz-0-74je9
events.out.tfevents.1528760355.mnist-simple-gpu-dist-worker-fziz-0-9mvsd
5 directories, 14 files
在 Kubernetes 部署 TensorBoard, 并且指定之前训练的 NAS 存储
- apiVersion: extensions/v1beta1
- kind: Deployment
- metadata:
- labels:
- app: tensorboard
- name: tensorboard
- spec:
- replicas: 1
- selector:
- matchLabels:
- app: tensorboard
- template:
- metadata:
- labels:
- app: tensorboard
- spec:
- volumes:
- name: kubeflow-dist-nas-mnist
- persistentVolumeClaim:
- claimName: kubeflow-dist-nas-mnist
- containers:
- name: tensorboard
- image: tensorflow/tensorflow:1.7.0
- imagePullPolicy: Always
- command:
- /usr/local/bin/tensorboard
- args:
- --logdir
- /training/tensorflow/logs
- volumeMounts:
- name: kubeflow-dist-nas-mnist
- mountPath: "/training"
- ports:
- containerPort: 6006
- protocol: TCP
- dnsPolicy: ClusterFirst
- restartPolicy: Always
将该模板保存到 tensorboard.yaml, 并且创建 tensorboard:
- kubectl create -f tensorboard.yaml
- deployment "tensorboard" created
TensorBoard 创建成功后, 通过 kubectl port-forward 命令进行访问
- PODNAME=$(kubectl get pod -l app=tensorboard -o jsonpath='{.items[0].metadata.name}')
- kubectl port-forward ${PODNAME} 6006:6006
通过 http://127.0.0.1:6006 登录 TensorBoard, 查看分布式训练的模型和效果:
总结
利用 tf-operator 可以解决分布式训练的问题, 简化数据科学家进行分布式训练工作. 同时使用 Tensorboard 查看训练效果, 再利用 NAS 或者 OSS 来存放数据和模型, 这样一方面有效的重用训练数据和保存实验结果, 另外一方面也是为模型预测的发布做准备. 如何把模型训练, 验证, 预测串联起来构成机器学习的工作流 (workflow), 也是 Kubeflow 的核心价值, 我们在后面的文章中也会进行介绍.
来源: http://blog.51cto.com/13876536/2149561