如何通过保留该维度的特定索引元素来减少张量的维度?

How to reduce dimension of a tensor with keeping specific indexed elements of that dimension?

维度的减少类似于reduce_max() 所做的,不同之处在于我想要该维度中元素的特定索引,而不是简单地选择最大的一个。例如,我有一个 2x3 张量 A = [[0,1,2],[2,2,0]]。如果我应用 tf.argmax(A),我得到索引张量 [1, 1, 0]。我怎样才能使用这个索引张量 [1, 1, 0] 来获得张量 tf.reduce_max(A, 0) = [2, 2, 2]?

我不直接使用 tf.reduce_max 的原因是我想使用不同的索引张量而不是 argmax 索引张量来减少维度或保留索引值而不是该维度的最大值。

您可以使用 tf.gather_nd 函数来执行此操作,但您需要将 [1, 1, 0] 索引张量转换为二维张量。

这里我假设索引张量是一个numpy数组(可以通过调用.numpy()方法将tensorflow张量转换成numpy数组

idx = np.array([1, 1, 0])

idx = np.c_[idx[:, np.newaxis], np.arange(len(idx))]

print(idx)

# Output:
# array([[1, 0],
#        [1, 1],
#        [0, 2]])

这意味着:在使用上述 tf.gather_nd

时选择 (row1, col0)、(row1, col1) 和 (row0, col2)
A = tf.Variable([[0, 1, 2], [2, 2, 0]])

tf.gather_nd(A, idx)

会给你预期的 [2, 2, 2] 张量。