分布式TensorFlow簡介
前一篇《分布式TensorFlow集群local server使用詳解》我們介紹了分布式TensorFlow的基本概念,現在我們可以動手搭建一個真正的分布式TensorFlow集群。
分布式TensorFlow集群由多個服務端進程和客戶端進程組成,在某些場景下,服務端和客戶端可以寫到同一個Python文件並起在同一個進程,但為了簡化代碼讓大家更好理解分布式架構,我們將啟動兩個worker並使用單獨的客戶端進程。
確認TensorFlow版本
首先我們需要安裝和確認TensorFlow的版本,注意0.8版本以前的TensorFlow不支持分布式,使用以前版本需要重新合Patch和打包。
python -c "import tensorflow; print(tensorflow.__version__)"
使用TensorFlow容器
如果本地已安裝Docker,通過容器使用TensorFlow環境更加簡單,只需一行命令。
sudo docker run -it tensorflow/tensorflow bash
Docker除了實現資源隔離,還可以管理不同版本的環境,例如可以很容易試用最新的RC版本。
sudo docker run -it tensorflow/tensorflow:r0.9rc0 bash
實現TensorFlow服務端
很多開發者讀過TensorFlow官方的Distributed Guide,想把分布式TensorFlow運行起來卻十分困難,主要原因是官方文檔沒有提供完整的例子,而且提供的代碼片段只能在0.9中運行,因此很多人嘗試修改那段代碼還是跑不起來。
其實分布式TensorFlow使用非常簡單,我們並不需要一個通用的程序,只要用幾行代碼分別實現服務端和客戶端即可,最簡單的服務端代碼如下worker1.py。
import tensorflow as tf
worker1 = "10.235.114.12:2222"
worker2 = "10.235.114.12:2223"
worker_hosts = [worker1, worker2]
cluster_spec = tf.train.ClusterSpec({ "worker": worker_hosts})
server = tf.train.Server(cluster_spec, job_name="worker", task_index=0)
server.join()
這里我們定義了兩個worker,其中job名都是“worker”,官方文檔中還定義了名為“ps”的job,實際上有沒有都可以,而worker可以是本地不同端口的兩個進程或者多台服務器上的進程。
為了模擬分布式環境,我們編寫worker2.py啟動第二個worker,注意代碼上稍有區別,因為index變了不能與前面的沖突,為了避免端口被容器隔離我們可以使用原來的Docker容器。
import tensorflow as tf
worker1 = "10.235.114.12:2222"
worker2 = "10.235.114.12:2223"
worker_hosts = [worker1, worker2]
cluster_spec = tf.train.ClusterSpec({ "worker": worker_hosts})
server = tf.train.Server(cluster_spec, job_name="worker", task_index=1)
server.join()
可以看到worker1和worker2分別監聽本地的2222和2223端口,接下來可以單獨寫客戶端應用,連接這兩個targe即可。
實現TensorFlow客戶端
Client的代碼我們寫得稍微復雜些,就是構造寫線性數據,其中斜率是2、截距是10,如果梯度下降算法正確的話最終w和b的輸出應該也接近2和10。
如我們所料,最終結果相當接近2和10,同時指定了第一個worker的CPU來執行梯度下降的算法,通過cluster spec我們還可以定義更靈活的集群,客戶端也可以通過"tf.device"來動態指定CPU和GPU計算資源。在生產環境下,官方推薦使用ps服務器存儲Variables,而ps其實是和我們定義的worker一樣,只是job名不同,架構可以參見官方文檔的tasks圖。
最后總結
希望通過這個完整的使用教程,大家都能親自實現分布式TensorFlow集群,並且編寫靈活的服務端和客戶端應用。
而在搭建過程中,大家可能發現TensorFlow只是一個深度學習的Library,我們需要實現和部署服務端、客戶端應用,而在集群定義中存在一個較嚴重的問題,就是cluster spec需要在進程啟動時指定,無法實現動態的擴容或縮容,這個問題社區希望通過引入Kubernetes集群管理工具來解決。還有一個問題就是我們的服務端應用啟動時只能bind localhost,我們也在Github建了對應Issue,如果大家都這些問題感興趣,也請繼續關注我和我們后續的文章。