Numpy 从Tensorflow 2.1 tensor中每批提取一个元素

假设我有一个由两个张量组成的批次，并且补丁中的张量大小为3

data = [[0.3, 0.5, 0.7], [-0.3, -0.5, -0.7]]

现在我想从面片中的每个张量中提取一个基于索引的元素：

index = [0, 2]

因此，输出应为

out = [0.3, -0.7] # Get index 0 from the first tensor in the batch and index 2 from the second tensor in the batch.

当然，这应该可以扩展到大批量。

索引的维度等于批量大小

我试图应用
tf.gather
和
tf.gather\n
但没有得到我想要的结果

例如，下面的代码打印上面指定的所需结果：

data = [[0.3, 0.5, 0.7], [-0.3, -0.5, 0.7]]

index = [0, 2]
out = tf.gather_nd(data, index)

print(out.numpy())

如果知道批量大小，可以执行以下操作：

import tensorflow as tf
data = tf.constant([[0.3, 0.5, 0.7], [-0.3, -0.5, 0.7]])

index = [0,2]
gather_inds = np.stack([np.arange(len(index)), index], axis=1)
out = tf.gather_nd(data, gather_inds)

你的聚集之所以不起作用，是因为你是从最内在的维度聚集。因此，您的索引需要与
数据的秩
张量相同。换句话说，你的指数应该是

[0,0] and [1,2]

作品非常感谢。