python - tf.nn.embedding_lookup 函数有什么作用？

Question

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

我无法理解这个函数的职责。它像查找表吗？即返回每个id对应的参数(在ids中)？

例如，在 skip-gram 模型中，如果我们使用 tf.nn.embedding_lookup(embeddings, train_inputs)，那么对于每个 train_input 它找到了对应的嵌入？

Answer 1

是的，这个功能很难理解，直到你明白这一点。

在最简单的形式中，它类似于 tf.gather。它根据ids指定的索引返回params的元素。

例如(假设你在 tf.InteractiveSession() 中)

params = tf.constant([10,20,30,40])
ids = tf.constant([0,1,2,3])
print tf.nn.embedding_lookup(params,ids).eval()

会返回[10 20 30 40]，因为params的第一个元素(索引0)是10，params的第二个元素(索引1)是20等

同样，

params = tf.constant([10,20,30,40])
ids = tf.constant([1,1,3])
print tf.nn.embedding_lookup(params,ids).eval()

将返回 [20 20 40]。

但 embedding_lookup 不止于此。 params 参数可以是张量的列表，而不是单个张量。

params1 = tf.constant([1,2])
params2 = tf.constant([10,20])
ids = tf.constant([2,0,2,1,2,3])
result = tf.nn.embedding_lookup([params1, params2], ids)

在这种情况下，ids 中指定的索引根据 分区策略 对应于张量的元素，其中默认分区策略是“mod”。

在 'mod' 策略中，索引 0 对应于列表中第一个张量的第一个元素。索引 1 对应于 second 张量的 first 元素。索引 2 对应于 third 张量的 first 元素，以此类推。只需索引 i 对应于第 (i+1) 个张量的第一个元素，对于所有索引 0..(n-1)，假设 params 是一个列表n 个张量。

现在，索引 n 不能对应张量 n+1，因为列表 params 仅包含 n 张量。所以索引 n 对应于第一个张量的 second 元素。同理，索引n+1对应第二张量的第二个元素，以此类推

所以，在代码中

params1 = tf.constant([1,2])
params2 = tf.constant([10,20])
ids = tf.constant([2,0,2,1,2,3])
result = tf.nn.embedding_lookup([params1, params2], ids)

索引0对应第一个张量的第一个元素:1

索引1对应第二张量的第一个元素:10

索引2对应第一个张量的第二个元素:2

索引3对应第二张量的第二个元素:20

因此，结果将是:

[ 2  1  2 10  2 20]