在TensorFlow中将张量的k个最大元素设置为零

Question

我想找到每行h的k个最大元素,并将零值设置为那些最大元素.

我可以使用top_k函数选择每行最高值的索引,如:

top_k = tf.nn.top_k(h, 1)

但我无法使用top_k返回的索引来更新张量.

我怎样才能做到这一点？提前致谢...

Answer 1

这有点棘手,也许有更好的解决方案.tf.scatter_update()在这里不起作用,因为它只能修改沿第一维的张量部分(例如,不是第一行和第二列中的元素).

你必须得到的values,并indices从tf.nn.top_k()创建稀疏张量,并减去初始张量x:

x = tf.constant([[6., 2., 0.], [0., 4., 5.]])  # of type tf.float32

k = 2
values, indices = tf.nn.top_k(x, k, sorted=False)  # indices will be [[0, 1], [1, 2]], values will be [[6., 2.], [4., 5.]]

# We need to create full indices like [[0, 0], [0, 1], [1, 2], [1, 1]]
my_range = tf.expand_dims(tf.range(0, indices.get_shape()[0]), 1)  # will be [[0], [1]]
my_range_repeated = tf.tile(my_range, [1, k])  # will be [[0, 0], [1, 1]]

# change shapes to [N, k, 1] and [N, k, 1], to concatenate into [N, k, 2]
full_indices = tf.concat([tf.expand_dims(my_range_repeated, 2), tf.expand_dims(indices, 2)], axis=2)
full_indices = tf.reshape(full_indices, [-1, 2])

to_substract = tf.sparse_to_dense(full_indices, x.get_shape(), tf.reshape(values, [-1]), default_value=0.)

res = x - to_substract  # res should be all 0.