前n個值的Tensorflow指標矩陣

Question

有誰知道如何提取排名2張量的每行的前n個最大值？

例如，如果我希望形狀[2,4]的張量的前2個值具有值：

[[40,30,20,10]，[10,20,30,40]]

所需的條件矩陣如下所示：[[True，True，False，False]，[False，False，True，True]]

一旦我有了條件矩陣，我就可以使用tf.select來選擇實際值。

謝謝你的幫助！

Answer 1

你可以使用內置的tf.nn.top_k函數來完成它：

a = tf.convert_to_tensor([[40, 30, 20, 10], [10, 20, 30, 40]])
b = tf.nn.top_k(a, 2)

print(sess.run(b))
TopKV2(values=array([[40, 30],
   [40, 30]], dtype=int32), indices=array([[0, 1],
   [3, 2]], dtype=int32))

print(sess.run(b).values))
array([[40, 30],
       [40, 30]], dtype=int32)

要獲得布爾值True/False值，您可以先獲取第k個值，然后使用tf.greater_equal ：

kth = tf.reduce_min(b.values)
top2 = tf.greater_equal(a, kth)
print(sess.run(top2))
array([[ True,  True, False, False],
       [False, False,  True,  True]], dtype=bool)

Answer 2

你也可以使用tf.contrib.framework.argsort

a = [[40, 30, 20, 10], [10, 20, 30, 40]]
idx = tf.contrib.framework.argsort(a, direction='DESCENDING')  # sorted indices
ranks = tf.contrib.framework.argsort(idx, direction='ASCENDING')  # ranks
b = ranks < 2  
# [[ True  True False False] [False False  True  True]]

此外，您可以使用1d張量替換2 ，以便每個行/列可以具有不同的n值。