tensorflow分布式训练ps负载不均衡问题

在使用euler0.1.2框架分布式训练GCN的工作中，遇到了这样一个问题：

• 在使用多个ps时，其中一个ps承载了大部分的参数，而其余的ps都在“划水”
• 当gcn的节点数量以及参数数量变多以后，ps的内存上限会成为训练的瓶颈，因为无法通过增加ps的数量，来放置增多的参数

好心的同事告诉了我这样一个方法：可以通过分割参数，来达到平均分配到每个ps的目的。具体实现，借助了tensorflow的两个接口：

• tf.fixed_size_paritioner
• tf.get_variable

参考博客：
https://blog.csdn.net/u013431916/article/details/80330813
https://moontree.github.io/2020/09/07/tf-embedding-lookup/

1.看下面的代码：

cluster = tf.train.ClusterSpec({
      'ps': flags_obj.ps_hosts,
      'worker': flags_obj.worker_hosts
  })
  server = tf.train.Server(
      cluster, job_name=flags_obj.job_name, task_index=flags_obj.task_index)

  if flags_obj.job_name == 'ps':
    server.join()
  elif flags_obj.job_name == 'worker':
    with tf.device(
        tf.train.replica_device_setter(
            worker_device='/job:worker/task:%d' % flags_obj.task_index,
            cluster=cluster):
     partitioner = tf.fixed_size_paritioner(num_ps, axix=0)
     embeddings = tf.get_variable(
        'embeddings',
        shape=[num, dim],
        initializer=tf.truncated_normal_initializer(),
        partitioner=partitioner)

遗憾的是，问题并没有得到解决

2.继续查询资料，发现tf.train.replica_device_setter这个函数主要实现变量在ps上初始化，其中提供了一个参数ps_strategy，即变量初始化的策略，这参数是可选的，不填的话默认循环去初始化。而官方提供的初始化策略还有两种：
2.1 随机初始化：tf.contrib.training.RandomStrategy
2.2 初始化时，每次选负载最小的ps分配参数（GLBS）：
tf.contrib.training.GreedyLoadBalancingStrategy

参考博客：https://blog.csdn.net/u012133034/article/details/81167040

采用GLBS策略，理论上是更合理的方式，遗憾的是问题还没解决

3.分析之后，发现之前是按照ps的数量还决定参数划分的份数，可能粒度还是太大，于是将分割的份数设为1000，至此，问题解决