在写ps分布式代码中,有一个很关键的问题,就是worker一旦训练完毕后,ps需要等待所有worker都停止,然后所有的ps再停止,而且需要紧密相连,根据博主的实践,为大家提供两种可行的办法,读者可以适当的选择合适自己的方法
1:第一个比较简单
def main(argv=None):
if FLAGS.job_name == 'ps':
with tf.device("/cpu:0"):
server.join()
然后worker不用做特殊的处理,只用正常结束就可以,ps就会得到worker的结束信息,但是博主在使用这个的时候遇到些问题,觉得这个还是不太好。博主当时用的是 tf.train.Supervisor 来作为所有worker和ps的管理器,也可能跟这个有关,但是后来尝试出来一种合适的解法,比较复杂,下面就写出来
2:
(1):ps和worker共享一个queue,然后ps刚开始就进行dequeue,因为队列是空的,所以会阻塞住。但是worker执行完的时候,会主动往这个队列中赛个数字,以此来触发ps的执行,这样ps和worker之间就可以进行通信了
(2):多个worker之间,一般是worker0最早结束,但是对于Supervisor来说,这样是不太完美的,可能会遇到一些问题,所以多个worker之间也需要进行通信,以保证可以同时结束,方法和上面一样,也是采用阻塞队列的方式
下面给出大概的写法
with tf.device("/job:worker/task:%d" % (i)):
return tf.FIFOQueue(FLAGS.num_worker, tf.int32, shared_name="worker_done_queue" + str(i))