如何在张量流中找到二维张量中的前k个值 [英] How to find the top k values in a 2-D tensor in tensorflow

问题描述

是否可以在Tensorflow的二维张量中找到顶部的k值?

Is there a way to find the top k values in a 2-D tensor in Tensorflow?

我可以将tf.nn.top_k用于一维张量，但不能用于二维张量.我有一个未知尺寸的二维张量，是否可以找到顶部的k值及其索引?

I can use tf.nn.top_k for a 1-D tensor but it cannot work with a 2-D tensor. I have a 2-D tensor with unknown size, is there a way to find the top k values and their indices?

非常感谢.

推荐答案

您可以在tf.nn.top_k()之前将矩阵重塑为一维张量，然后从一维索引中计算二维索引:

You can reshape your matrix to a 1-D tensor before tf.nn.top_k(), then compute the 2-D indices from the 1-D ones:

x = tf.random_uniform((3, 4))
x_shape = tf.shape(x)
k = 3

top_values, top_indices = tf.nn.top_k(tf.reshape(x, (-1,)), k)
top_indices = tf.stack(((top_indices // x_shape[1]), (top_indices % x_shape[1])), -1)

with tf.Session() as sess:
    mat, val, ind = sess.run([x, top_values, top_indices])
    print(mat)
    # [[ 0.2154634   0.52707899  0.29711092  0.74310601]
    #  [ 0.61274767  0.82408106  0.27242708  0.25479805]
    #  [ 0.25863791  0.16790807  0.95585966  0.51889324]]
    print(val)
    # [ 0.95585966  0.82408106  0.74310601]
    print(ind)
    # [[2 2]
    #  [1 1]
    #  [0 3]]

这篇关于如何在张量流中找到二维张量中的前k个值的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！