介紹

本系列將介紹如何在阿里雲容器服務上執行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擅長的，固化和可重複的操作儲存到容器映象。而自動化的構建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對應的格式如下所示：

{  
    "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訓練所需要的資料

原文連結