How to scale tf.nn.embedding_lookup_sparse

问题

I'm trying to build a very large sparse model (e.g. LR if there is only one embedding layer), the input dimension can be as large as 100000000, and the sample is very sparse, the average number of non zero value is around 100. Since the weights is very large and we have to partition and distribute it onto different servers. Here is the code:

weights = tf.get_variable("weights",                                        
                          weights_shape,                                    
                          partitioner=tf.fixed_size_partitioner(num_shards, axis=0), 
                          initializer=tf.truncated_normal_initializer(stddev=0.1))
biases = tf.get_variable("biases",                                          
                         biases_shape,                                      
                         initializer=tf.truncated_normal_initializer(stddev=0.1))

result = tf.nn.embedding_lookup_sparse(weights,                              
                                       ids,                                  
                                       values,                               
                                       partition_strategy="div",             
                                       combiner="sum") + biases

This is the generated graph for this op

From the graph, embedding_lookup_sparse just simply cancat the sharded weights, which cause a huge of unnecessary network traffic. This looks stupid. The resonnable way is making lookup for each shard in local indepently, and just send back the lookuped results and aggregate them. In this way, the traffic is dramatically reduced. I'm wondering whether TensorFlow support this mode? Of course, I can do this by customized code.

Edited contents

The solution that works as expected:

num_shards = 2
weights = []
assert weights_shape[0] % num_shards == 0
for i in range(0, num_shards):
  weights_i = tf.get_variable("weights-%02d" % i,
                              [weights_shape[0]/num_shards] + weights_shape[1:],
                              initializer=tf.truncated_normal_initializer(stddev=0.1))
  weights.append(weights_i)                                                 
biases = tf.get_variable("biases",
                         biases_shape,
                         initializer=tf.truncated_normal_initializer(stddev=0.1))

result = tf.nn.embedding_lookup_sparse(weights,                              
                                       ids,                                  
                                       values,                               
                                       partition_strategy="div",             
                                       combiner="sum") + biases

来源：https://stackoverflow.com/questions/40628977/how-to-scale-tf-nn-embedding-lookup-sparse

标签

tensorflow

sharding

partition