Tensorflow:每行索引 [英] Tensorflow: index per row

问题描述

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

Suppose I have a Tensor of shape (100,20). Now I also have a Tensor of indices of shape (100,). How to obtain now a Tensor of shape (100,) or (100,1) with per row (100 rows) the right value (selected by the corresponding index in indices?

小例子:所以让我们说张量 A 是

Small example: So let's say tensor A is

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

张量 B 是

[0,2,1]

然后我想要作为输出

[1,6,8]

推荐答案

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

You can join your B tensor with an appropriate range to create two-dimensional indices (in your example [[0, 0], [1, 2], [2, 1]]) and then extract the elements using 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)

这篇关于Tensorflow:每行索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！