python - 使用 Tensorflow 的 top_k 和 scatter_nd

Question

我正在尝试在 tensorflow 中编写一个仅传播每个特征图的前 k 个值的操作。

例子:

k=1，输入大小是[batch_size, x, y, channels] 假设它是[1,2,2,3]

输出应该是相同的大小，如果 k=1 则每个 x,y 平面将只有一个非零值。

numpy 中的示例:

input = [[[[6.4 1.4 1.3] [2.1  6.5  4.8]][[2.3 9.2  2.8][7.9  5.1 0.6]]]]]

输出应该是:

[[[[6.4 0. 0.] [0. 6.5 0.]]  [[0. 9.2 0.] [7.9 0. 0.]]]]

为了在 tensorflow 中执行此操作，我想使用 nn.top_k，然后使用 scatter_nd。

问题是 top_k 返回请求元素的索引与 scatter_nd 需要它的方式非常不同。

top_k 以 (1,2,2, 1)

scatter_nd 需要它作为每个值的所有坐标的列表，如下所示:

[[0, 0, 0, 0], [0, 0, 1, 1], [0, 1, 0, 1], [0, 1, 1, 0]]

有谁知道在它们之间进行转换的方法吗？甚至可能完全针对此操作采取不同的方法？

Answer 1

tf.nn.top_k()仅返回最后一个维度 中的前 k 个值。所以你必须加回所有其他维度。最简单的 tf.where() .代码(测试):

import tensorflow as tf

inp = tf.constant( [ [ [ [6.4, 1.4, 1.3], [2.1,  6.5,  4.8] ], [ [2.3, 9.2, 2.8], [7.9, 5.1, 10.6] ] ] ] )

t, idx = tf.nn.top_k( inp, k = 2 )
idx_one_hot = tf.one_hot( idx, depth = 3 )
idx_red = tf.reduce_sum( idx_one_hot, axis = -2 )
idx2 = tf.where( tf.not_equal( idx_red, 0 ) )

with tf.Session() as sess:
    print( sess.run( idx2 ) )

输出(请注意，我已将您示例中的最后一个数字更改为索引也为 2，只有 0 和 1 看起来有点误导，好像它是 bool 张量):

[[0 0 0 0]
[0 0 0 1]
[0 0 1 1]
[0 0 1 2]
[0 1 0 1]
[0 1 0 2]
[0 1 1 0]
[0 1 1 2]]

请注意，这会丢失 top_k 报告的最后一个维度中的索引顺序，它会将其更改为索引本身的递增顺序。