Как найти верхние значения k в двумерном тензоре в тензорном потоке

Question

Есть ли способ найти верхние значения k в 2-D тензоре в Tensorflow?

Я могу использовать tf.nn.top_k для 1-D тензора, но он не может работать с 2-D тензор.У меня есть двумерный тензор с неизвестным размером, есть ли способ найти верхние значения k и их индексы?

Большое спасибо.

Answer 1

Вы можете изменить свою матрицу до 1-D тензора до tf.nn.top_k(), а затем вычислить 2-D индексы из 1-D:

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]]

Answer 2

Один из способов сделать это - изменить форму, например, xx= np.reshape(x,(-1,)), а затем что-то вроде x[:k] подойдет?