在tensorflow的分布式训练中,默认将参数以变量名的维度分布式存储在不同ps上。如将item_embedding存储在ps0上,item_bias存储在ps1上。默认不会将同一个变量存储在不同的ps上。
在实际情况中,有的变量行数很多,如item_embedding的行数一般有几百万量级,如果存储在同一个ps上,会导致该ps成为性能瓶颈。为了解决以上问题,在tensorflow中可以通过定义partitioner的方式将同一个变量分布式存储在不同ps上。具体使用方法如下:
self.partitioner = tf.fixed_size_partitioner(10) #定义partitioner,均匀分布存储在10个ps上
with tf.variable_scope('embedding',partitioner=self.partitioner): #定义variable时加上partitioner
self.item_embedding = tf.get_variable('item_embedding',[item_num, emb_size])
以上定义会将item_embedding均匀分布存储在10个ps上,如下图所示:
注意在直接访问item_embedding的地方,都要加上with tf.variable_scope('xxxx',partitioner=self.partitioner):,如下面的语句
with tf.variable_scope("network_part",partitioner=self.partitioner):
item_list_emb = tf.nn.embedding_lookup(self.item_embedding,features['itemlist'])
后面继续访问item_list_emb,则可以不需要加上with tf.variable_scope('xxxx',partitioner=self.partitioner):了
另外在导出模型中item_embedding时,需要在变量名后加上/part_xx,如embedding/item_embedding/part_0。
如需导出所有item_embedding,将变量名embedding/item_embedding/part_0,embedding/item_embedding/part_1,.... embedding/item_embedding/part_9分别导出后合并到一个文件里。