如何在 Tensorflow 张量中仅舍入 k 个最大元素

Question

假设有一个 TensorFlow 张量 - 例如 [0.1,0.2,0.3,0.4]。我想将 k 个最大的元素四舍五入，然后将其余的四舍五入。 （例如，当k=2时，我希望得到[0,0,1,1]。当k=3时，我希望得到[0,1 ,1,1].)

我想仅使用 TensorFlow 操作来实现此功能。我如何做到这一点？

Answer 1

试试这样的：

import tensorflow as tf

x = tf.constant([0.1,0.2,0.3,0.4])
k = 3
greatest = tf.math.top_k(x, k=k).indices 
tensor = tf.tensor_scatter_nd_update(tf.zeros_like(x), tf.reshape(greatest, (tf.shape(greatest)[0], 1)), tf.ones_like(tf.gather(x, greatest)))

k = 3：

tf.Tensor([0. 1. 1. 1.], shape=(4,), dtype=float32)

k = 2：

tf.Tensor([0. 0. 1. 1.], shape=(4,), dtype=float32)

此方法不会真正舍入，因为将 0.3 和 0.4 舍入到最接近的整数会导致零，这不是您想要的。所以我只是将张量中最高的 k 值转换为 1，其余的值转换为零，但如果它仍然是二进制分类，这对于您的用例来说应该足够了。

如果您真的想向上取整最大的k 值，请使用tf.math.ceil 而不是tf.ones_like：

tensor = tf.tensor_scatter_nd_update(tf.zeros_like(x), tf.reshape(greatest, (tf.shape(greatest)[0], 1)), tf.ceil((tf.gather(x, greatest))))

Answer 2

您可以为此使用tf.math.top_k。该函数将返回给定张量中 k 最大元素的值和索引。

https://www.tensorflow.org/api_docs/python/tf/math/top_k

然后您可以使用返回的索引，将张量中的值设置为特定值。

以下解决方案对问题中提到的值进行四舍五入。

import tensorflow as tf

x = tf.constant([0.1,0.2,0.3,0.4])
k = 3

# retrieve min and max values
max_value = tf.math.ceil(tf.math.reduce_max(x))
min_value = tf.math.floor(tf.math.reduce_min(x))

# retrieve the k largest elements
k_largest = tf.math.top_k(x, k=k)

# reshape the indices, required for ‘scatter‘ function
indices = tf.reshape(k_largest.indices, (-1,1))
values = k_largest.values

# initialize update tensor with max_value
updates = max_value * tf.ones_like(values)
# initialize result with min_value
x_new = min_value * tf.ones_like(x)
# update values for k_largest indices
x_new = tf.tensor_scatter_nd_update(
    x_new, indices, updates)

print(x_new)

如果您要求的ceil 和floor 操作应应用于每个元素，而不是应用于张量内的min 和max 值，则如下所示：

import tensorflow as tf

x = tf.constant([0.1,0.2,0.3,0.4])
k = 3

# retrieve the k largest elements
k_largest = tf.math.top_k(x, k=k)
# reshape the indices, required for ‘scatter‘ function
indices = tf.reshape(k_largest.indices, (-1,1))

# get floored values
floored_values = tf.math.floor(x)
# get ceiled values only for top-k
ceiled_values = tf.math.ceil(k_largest.values)

# initialize result with per element floored values
x_new = floored_values
# update values for k_largest indices with per element ceiled values
x_new = tf.tensor_scatter_nd_update(
    floored_values, indices, ceiled_values)

print(x_new)