有谁知道如何提取等级2张量的每行的前n个最大值?
例如,如果我想要一个形状张量[2,4]的前2个值,其值为:
[[40,30,20,10],[10,20,30,40]]
所需的条件矩阵如下所示: [[True,True,False,False],[False,False,True,True]]
一旦我有条件矩阵,我就可以使用tf.select来选择实际值。
感谢您的帮助!
答案 0 :(得分:9)
您可以使用内置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)
答案 1 :(得分:0)
您还可以使用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]]
此外,您可以将2
替换为一维张量,以便每个行/列可以具有不同的n
值。