几个小时几行代码简单直白 tensorflow 分布式代码实战

骑猪看日落 2023-07-03 10:47 36阅读 0赞

北京又下雪了，这个2020年开启的方式太沉重了。

**注意，市面上普遍的tf分布式代码，我实测都有一个问题，如果你租用GPU资源请务必警惕这个问题会让你白花钱。如果你用的自己的资源，那应该是没有太大问题的。具体请看（**[https://blog.csdn.net/u013249853/article/details/104385793][https_blog.csdn.net_u013249853_article_details_104385793]**）**

**好烦，每次发文都要审核，链接就得等他审核完了上，好麻烦，看来我得另外找个网站写博客了。**

我本人对tf并不精通。分布式听起来也很复杂。而且我要做的是**将deeplabv3+的代码做成分布式的**。也就是多个节点，相当于用两台服务器的gpu。**听起来很麻烦，实际上居然很简单。而我自己也是花了几个小时从什么都不懂到代码跑起来。**

**这个教程目的是以tensorflow为例，将分布式细节写出来，补充一下网上教程没有的东西。**

很多教程写了特别长的文字，我认为画图更简单。**本教程尽量简短，好懂，并且有代码。只是基础教程。**

# **理论（特短，够用）** #

**一**台电脑=服务器=server是一个节点，包含了**多**个GPU。首先分布式的方式就是让**多**台电脑上的gpu共同干活。

分布式工作分为两个部分，parameter server（ps）以及worker。**眼熟ps与worker**，因为这个是工作，每个server，都得干活，所以只能是从这两个工作里面选择。ps的工作类似于存储参数，而损失的计算，梯度的决定都是有worker进行的。**这个对代码的影响就是，ps节点其实完全可以由cpu来做。worker必须由gpu做。**

整体结构如图：一共四个sever，每个sever假设包含4个GPU，下图一共16个GPU。两个server工作是ps，两个sever的工作是worker，这个name其实没有在代码中配置，所以不用理会。server同做一个工作，也需要区分的，所以又引入了task，并且有task id。**这里只是演示一下job(ps,worker)和server（节点）的关系。**![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70][]

# 代码 #

代码的讲解是踩点来的。就是怎么用代码互相交流。

从理论上看，**我们需要一些节点，并且给他们分配工作。**

所以做一下程序入口接受参数（节点都是是谁，给什么工作了），我比较喜欢接收参数，不喜欢在代码里面写死。因为flags是tf基础，不想解释增加长度。

每个节点都得被单独通知，并且单独运行，这意味着如果你有一个ps,两个worker（一般用一个ps即可），你得在bash命令里：

*python train.py --ps都谁(ps\_hosts) --worker都谁(woker\_hosts) --我被分配干啥(job\_name) --我是第几个干这活的(task\_index)*

**就是输入三次，跑三次，同时。ps和worker都会等你输完了在一起工作，毕竟要等同伴。**

下面是**一个节点**需要知道的消息。

flags = tf.app.flags
    flags.DEFINE_integer('task_index', 0, 'The task ID.0 then is cheif session')
    flags.DEFINE_string("ps_hosts", "g0101:2222", "ps hosts")
    flags.DEFINE_string("worker_hosts", "g0101:2223,g0102:2224", "worker hosts")
    flags.DEFINE_string("job_name", "worker", "'ps' or'worker'")

这四个东西怎么用需要看下面的代码，不要着急。不过我们可以看出，上面有一个ps server（节点）g0101:2222（可以当做该节点的名字）,以及两个worker节点，g0101:2223+g0102:2224. 这里需要插播一下**节点工作的时候是知道都有哪些节点一起工作**，并且大家都是做什么任务的。

传参的时候是这样写的，这个节点需要知道的消息：

下面分别写下ps工作节点与worker工作节点传参。

**ps:**

python "${WORK_DIR}"/train_cloud.py \  
    --ps_hosts="g0101:2222" \
      --worker_hosts="g0101:2223,g0102:2224" \
      --job_name="ps" \
      --task_index="0"

然后节点就知道**自己被分配了ps工作，只有一个ps节点（ps\_hosts就一个，逗号是分隔符）。自己就是第一个。还有两个worker给自己打工**。

**worker1:**

python "${WORK_DIR}"/train_cloud.py \  
    --ps_hosts="g0101:2222" \
      --worker_hosts="g0101:2223,g0102:2224" \
      --job_name="worker" \
      --task_index="1"

然后节点就知道自己被分配了worker工作，并且知道有两个worker，自己是第二个（task\_index="1",计数从零开始，所以是第二个）。并且知道ps节点是谁。worker0当然就是类比着来。

python "${WORK_DIR}"/train_cloud.py \  
    --ps_hosts="g0101:2222" \
      --worker_hosts="g0101:2223,g0102:2224" \
      --job_name="worker" \
      --task_index="0"

**这里注意，我写的ps单独用了一个节点，如果想节省一下gpu，那么就写：**

**ps:**

CUDA_VISIBLE_DEVICES=""
    python "${WORK_DIR}"/train_cloud.py \  
    --ps_hosts="g0101:2223" \
      --worker_hosts="g0101:2223,g0102:2224" \
      --job_name="ps" \
      --task_index="0"

**worker1:**

python "${WORK_DIR}"/train_cloud.py \  
    --ps_hosts="g0101:2223" \
      --worker_hosts="g0101:2223,g0102:2224" \
      --job_name="worker" \
      --task_index="1"

**注意看**CUDA\_VISIBLE\_DEVICES，这时ps用的是worker1的CPU了。

然后需要建立一个cluster。需要用的是

ps_hosts=FLAGS.ps_hosts.split(",")
      worker_hosts=FLAGS.worker_hosts.split(",")
      cluster=tf.train.ClusterSpec({"ps":ps_hosts,"worker":worker_hosts})
      server=tf.train.Server(cluster,job_name=FLAGS.job_name,
                             task_index=FLAGS.task_index)
      if FLAGS.job_name == "ps":
          server.join()
      elif FLAGS.job_name =="worker":
          graph = tf.Graph()
          with graph.as_default():
            with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % (FLAGS.task_index),
                                                          cluster=cluster)):

with下面就是正常的dataset的设置，config的设置，并不是分布式独有的，所以省略了。而且是deeplab代码（改完可跑），所以建议看tf官方[简单实例][Link 1]，一定**看完哦，根本不长。**

然后在session 里面的参数：master和一般设置不一样。

with tf.train.MonitoredTrainingSession(
                    master=server.target,

# 运行 #

输入的命令大家已经看到了，的确是需要输三遍bash 命令运行同一个python，但是由于分配的工作不同，所以是不一样的。

就是这个，还差两个worker。

python "${WORK_DIR}"/train_cloud.py \  
    --ps_hosts="g0101:2222" \
      --worker_hosts="g0101:2223,g0102:2224" \
      --job_name="ps" \
      --task_index="0"

然后节点具体叫什么，这个需要问你的服务器管理员。你用我这个肯定不好使。

batchsize等的设置，每个节点，都要单独配置参数，比如worker0包含了一个GPU，worker1包含了两个GPU，那你的batchsize肯定是worker1更多。这个时候就把两个节点视作独立的。因为他们只能看到自己内部包含了什么设备（GPU），他虽然知道有个worker0和自己一起工作，却不知道worker0内部有什么设备。

# 输出与检验 #

虽然运行成功，但是怎么验证是否一起工作了呢？万一是各做各的就完蛋了。所以我又用deeplabv3+做了验证，如果batchsize不足的话，batch normalization的存在会使结果难看。

所以其输出是有三个ps0.out,worker0.out,worker1.out

我截一下图：

ps0,我用的就是cpu

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70 1][]

worker0

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70 2][]

worker1：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70 3][]

因为其实我有四个worker，这里只贴了两个。

参考：

[https://www.youtube.com/embed/la\_M6bCV91M][https_www.youtube.com_embed_la_M6bCV91M]

[https://github.com/tensorflow/examples/blob/master/community/en/docs/deploy/distributed.md][Link 1]

[https://henning.kropponline.de/2017/03/19/distributing-tensorflow/][https_henning.kropponline.de_2017_03_19_distributing-tensorflow]

[https://zhuanlan.zhihu.com/p/35083779][https_zhuanlan.zhihu.com_p_35083779]

[https://www.jianshu.com/p/fdb93e44a8cc][https_www.jianshu.com_p_fdb93e44a8cc]

[https_blog.csdn.net_u013249853_article_details_104385793]: https://blog.csdn.net/u013249853/article/details/104385793
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20200205192823272.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM=,size_16,color_FFFFFF,t_70
[Link 1]: https://github.com/tensorflow/examples/blob/master/community/en/docs/deploy/distributed.md
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70 1]: https://img-blog.csdnimg.cn/20200212013157963.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM=,size_16,color_FFFFFF,t_70
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70 2]: https://img-blog.csdnimg.cn/20200212013258276.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM=,size_16,color_FFFFFF,t_70
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM_size_16_color_FFFFFF_t_70 3]: https://img-blog.csdnimg.cn/20200212013327670.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3UwMTMyNDk4NTM=,size_16,color_FFFFFF,t_70
[https_www.youtube.com_embed_la_M6bCV91M]: https://www.youtube.com/embed/la_M6bCV91M
[https_henning.kropponline.de_2017_03_19_distributing-tensorflow]: https://henning.kropponline.de/2017/03/19/distributing-tensorflow/
[https_zhuanlan.zhihu.com_p_35083779]: https://zhuanlan.zhihu.com/p/35083779
[https_www.jianshu.com_p_fdb93e44a8cc]: https://www.jianshu.com/p/fdb93e44a8cc