tensorflow分布式训练ps负载不均衡问题
2020-10-15 本文已影响0人
钢镚儿_e134
在使用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,至此,问题解决