Question

考虑情况：

token_ids = [17, 189, 981, 1000, 11, 42, 109, 26, 3377, 261]  
word_ids = [0, 0, 0, 0, 1, 1, 1, 2, 2, 2]

我需要像这样计算每个token_ids减少的word_id的总和：

output = [ (emb[17] + emb[189] + emb[981] + emb [1000]),  
           (emb[11] + emb[42] + emb[109]),
           (emb[26] + emb[3377] + emb[261]) ]

其中emb是任何嵌入矩阵。

我可以像这样使用for循环在python中编写此代码：

prev = 0
sum_all = []
sum = 0
for i in range(len(word_ids)):
    if word_ids[i] == prev:
        sum += emb[token_ids[i]]
    else:
        sum_all += [sum]
        sum = emb[token_ids[i]]
        prev = word_ids[i]
    if i == len(word_ids):
        sum_all += [sum]
return sum_all

但是我想在 tensorflow 中有效地做到这一点（如果可能的话，向量化）。有人可以提出建议怎么做吗？

Answer 1

您需要tf.segment_sum来计算张量段上的总和。

import tensorflow as tf

token_ids = tf.constant([17, 189, 981, 1000, 11, 42, 109, 26, 3377, 261],tf.int32)
word_ids = tf.constant([0, 0, 0, 0, 1, 1, 1, 2, 2, 2],tf.int32)

emb_matrix = tf.ones(shape=(4000,3))
emb = tf.nn.embedding_lookup(emb_matrix, token_ids)

result = tf.segment_sum(emb,word_ids)

with tf.Session() as sess:
    print(sess.run(result))

[[4. 4. 4.]
 [3. 3. 3.]
 [3. 3. 3.]]

张量流中减少总数的可变数量

1 个答案: