python - 如何为 4d 张量中的 k 个最大元素创建一个单热张量？

Question

给定一个形状张量(2,3,3,1)，其中批量大小为 2，每个矩阵的形状为(3,3,1)。如何从每个矩阵中找到 k 个最大元素并创建一个 one-hot 矩阵，使得这 k 个位置的条目为 1，其他位置为 0。示例:(请注意，为了简单起见，每个条目都是 float ，使用整数)

input_tensor=[[[1, 5, 7],
      [2, 8, 1],
      [3, 9, 1],
     ],
     [[0, 9, 5],
      [6, 0, 4],
      [3, 0, 8]
     ]
    ]

k=3 的单热张量:

output_tensor=[[[0, 0, 1],
      [0, 1, 0],
      [0, 1, 0],
     ],
     [[0, 1, 0],
      [1, 0, 0],
      [0, 0, 1]
     ]
    ]

tf.nn.top_k 将仅返回最后一个维度中的 k 个最大元素。如何从 3d 张量中获取 k 个最大元素 ex:(3,3,1)。此外 tf.one_hot 会在指定深度和给定索引的每一行中放置一个 1，但此处情况并非如此。

Answer 1

如果你想从每个矩阵中找到k个最大的元素，可以使用以下方法。

import tensorflow as tf

input_tensor = tf.constant([[[1, 5, 7],[2, 8, 1],[3, 9, 1]],
                            [[0, 9, 5],[6, 0, 4],[3, 0, 8]]],dtype=tf.int32)

k_tf = tf.placeholder(shape=(),dtype=tf.int32)

temp = tf.reshape(input_tensor,shape=(input_tensor.shape[0],-1))
# [[1 5 7 2 8 1 3 9 1]
#  [0 9 5 6 0 4 3 0 8]]
result = tf.reduce_sum(tf.one_hot(indices=tf.nn.top_k(temp,k=k_tf)[1], depth=temp.shape[1]), axis=1)
# [[0. 0. 1. 0. 1. 0. 0. 1. 0.]
#  [0. 1. 0. 1. 0. 0. 0. 0. 1.]]
result = tf.reshape(result,input_tensor.shape)

with tf.Session() as sess:
    print('k=2:')
    print(sess.run(result, feed_dict={k_tf: 2}))
    print('k=3:')
    print(sess.run(result,feed_dict={k_tf:3}))

k=2:
[[[0. 0. 0.]
  [0. 1. 0.]
  [0. 1. 0.]]

 [[0. 1. 0.]
  [0. 0. 0.]
  [0. 0. 1.]]]
k=3:
[[[0. 0. 1.]
  [0. 1. 0.]
  [0. 1. 0.]]

 [[0. 1. 0.]
  [1. 0. 0.]
  [0. 0. 1.]]]