Understand tf.nn.embedding_lookup(): Pick Up Elements by Ids – TensorFlow Tutorial

TensorFlow tf.nn.embedding_lookup() function can allow us to pick up elements by ids from a tensor. In this tutorial, we will introduce how to use this function in our application correctly.

Syntax

tf.nn.embedding_lookup(
    params,
    ids,
    partition_strategy='mod',
    name=None,
    validate_indices=True,
    max_norm=None
)

Pick up elements by ids from a tensor and return a new tensor.

Parameters

params: a tensor or numpy ndarray.

ids: the postion list of elements in params.

Return

Return a new tensor with positionss of elements in params.

Here is an example to show how to use this function.

import tensorflow as tf;
import numpy as np;
 
vec = np.random.random([10,5])
b = tf.nn.embedding_lookup(vec, [1, 3, 1])

init = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()
 
with tf.Session() as sess:
    sess.run([init, init_local])
    print('vec = ')
    print(vec)
    print("the look up elements by ids")
    print(sess.run(b))

In this example, we create a 10 * 5 matrix named as vec randomly, it may be:

vec = 
[[ 0.23774171  0.46927013  0.30450351  0.49490524  0.92624406]
 [ 0.21206469  0.75452645  0.27273611  0.89547634  0.10094618]
 [ 0.18285835  0.53391567  0.91681091  0.18240941  0.65784034]
 [ 0.6411879   0.76711249  0.04965164  0.85007949  0.11408826]
 [ 0.67376102  0.51645354  0.39302517  0.78245695  0.4591333 ]
 [ 0.66452224  0.92411268  0.0083408   0.14662294  0.47980183]
 [ 0.62246077  0.09974218  0.75115197  0.17576755  0.58825293]
 [ 0.39409025  0.00463679  0.70296569  0.9849676   0.40180546]
 [ 0.03795578  0.86115679  0.90238849  0.1182467   0.31533444]
 [ 0.04804927  0.01061292  0.81253765  0.7933884   0.78598149]]

The demension of vec is 2.

Then we want to pick up elements with postion index is [1, 3, 1] from vec to create a new tensor. We can do like:

b = tf.nn.embedding_lookup(vec, [1, 3, 1])

Then we can find b may be:

the look up elements by ids
[[ 0.21206469  0.75452645  0.27273611  0.89547634  0.10094618]
 [ 0.6411879   0.76711249  0.04965164  0.85007949  0.11408826]
 [ 0.21206469  0.75452645  0.27273611  0.89547634  0.10094618]]

From the b we can find:

1.The dimension of b is

D_b = D_ids + D_vec – 1

In this example,  D_ids = 1, D_vec = 2, then D_b = 1+2-1 = 2.

If D_ids = 2

b = tf.nn.embedding_lookup(vec, [[1], [3], [1]])

The b may be:

[[[ 0.8532426   0.03158583  0.11454635  0.3275151   0.9703617 ]]

 [[ 0.93875557  0.87250653  0.46802411  0.18365519  0.5412155 ]]

 [[ 0.8532426   0.03158583  0.11454635  0.3275151   0.9703617 ]]]

D_b = 2+2-1 = 3.

2.The count of elements in b is the same as to the count of elements in ids, which is 3 in this example.

3.The order of elements in b is the same as the order of elements in ids, which is [ 1, 3, 1].