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
