Tensorflow:每行索引

Question

假设我有一个形状为 (100,20) 的张量。现在我还有一个形状指数张量 (100,)。现在如何获得形状为 (100,) 或 (100,1) 且每行(100 行)正确值(由索引中的相应索引选择？

小例子:所以假设张量 A 是

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

张量B是

[0,2,1]

然后我想要作为输出

[1,6,8]

Answer 1

您可以将 B 张量加入适当的范围以创建二维索引(在您的示例中 [[0, 0], [1, 2], [2, 1]])然后使用 tf.gather_nd 提取元素:

b_2 = tf.expand_dims(b, 1)
range = tf.expand_dims(tf.range(tf.shape(b)[0]), 1)
ind = tf.concat(1, [range, b_2])
res = tf.gather_nd(a, ind)