How to use tf.nn.embedding_lookup_sparse in TensorFlow?

tobe picture tobe · Aug 29, 2016 · Viewed 8k times · Source

We have tried using tf.nn.embedding_lookup and it works. But it needs dense input data and now we need tf.nn.embedding_lookup_sparse for sparse input.

I have written the following code but get some errors.

import tensorflow as tf
import numpy as np

example1 = tf.SparseTensor(indices=[[4], [7]], values=[1, 1], shape=[10])
example2 = tf.SparseTensor(indices=[[3], [6], [9]], values=[1, 1, 1], shape=[10])

vocabulary_size = 10
embedding_size = 1
var = np.array([0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0])
#embeddings = tf.Variable(tf.ones([vocabulary_size, embedding_size]))
embeddings = tf.Variable(var)

embed = tf.nn.embedding_lookup_sparse(embeddings, example2, None)

with tf.Session() as sess:
    sess.run(tf.initialize_all_variables())

    print(sess.run(embed))

The error log looks like this.

enter image description here

Now I have no idea how to fix and use this method correctly. Any comment could be appreciated.

After diving into safe_embedding_lookup_sparse's unit test, I'm more confused why I got this result if giving the sparse weights, especially why we got something like embedding_weights[0][3] where 3 is not appeared in the code above.

enter image description here

Answer

rvinas picture rvinas · Aug 29, 2016

tf.nn.embedding_lookup_sparse() uses Segmentation to combine embeddings, which requires indices from SparseTensor to start at 0 and to be increasing by 1. That's why you get this error.

Instead of boolean values, your sparse tensor needs to hold only the indices of every row that you want to retrieve from embeddings. Here's your tweaked code:

import tensorflow as tf
import numpy as np

example = tf.SparseTensor(indices=[[0], [1], [2]], values=[3, 6, 9], dense_shape=[3])

vocabulary_size = 10
embedding_size = 1
var = np.array([0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0])
embeddings = tf.Variable(var)

embed = tf.nn.embedding_lookup_sparse(embeddings, example, None)

with tf.Session() as sess:
    sess.run(tf.initialize_all_variables())
    print(sess.run(embed)) # prints [  9.  36.  81.]

In addition, you can use indices from tf.SparseTensor() to combine word embeddings using one of the allowed tf.nn.embedding_lookup_sparse() combiners:

  • "sum" computes the weighted sum of the embedding results for each row.
  • "mean" is the weighted sum divided by the total weight.
  • "sqrtn" is the weighted sum divided by the square root of the sum of the squares of the weights.

For example:

example = tf.SparseTensor(indices=[[0], [0]], values=[1, 2], dense_shape=[2])
...
embed = tf.nn.embedding_lookup_sparse(embeddings, example, None, combiner='sum')
...
print(sess.run(embed)) # prints [ 5.]