Search code examples
pythontensorflowsparse-matrixembedding

How to use tf.nn.embedding_lookup_sparse in TensorFlow?


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


Solution

  • 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.]