在tensorflow中创建顶级“k”元素掩码的更有效方法

Question

我想创建一个函数，对于给定的1d，-tensor输出掩码，在位置上，对应于最上面的k值，在其他地方有1和0。也就是说，举例来说，我有：

tensor = [1, 0, 7, 5, 2, 3] : get_largest_mask(tensor, 3) = [0, 0, 1, 1, 0, 1]

我创建了以下函数：

def get_largest_mask(tensor, n_to_keep):
    # tensor 1-d tensor
    values, indices = tf.math.top_k(tensor, k=n_to_keep)

    mask = tf.zeros(tf.size(tensor))
    mask = tf.tensor_scatter_nd_update(mask, [[idx] for idx in indices], tf.ones(n_to_keep))

    return mask

然而，对于感兴趣的情况，它的工作速度相当慢，正如我测量的那样，大多数时间都是由tf.tensor_scatter_nd_update主导的。更快的替代方案是什么？

张量的典型大小是10^3-10^4元素，k的量级是‘10^2-10^3’。

Answer 1

我会找出最大的K值，然后进行大小比较

import tensorflow as tf

tensor = tf.convert_to_tensor([1, 0, 7, 5, 2, 3])

mask = tf.cast(tensor >= tf.math.top_k(tensor, 3)[0][-1], tf.int32)
# mask = <tf.Tensor: shape=(6,), dtype=int32, numpy=array([0, 0, 1, 1, 0, 1], dtype=int32)>

解释

tf.math.top_k返回两个值，第一个是包含实际前k个值的张量，第二个是索引。我们获取这些值，然后访问最小值的[-1]。然后，我们通过询问>=问题来创建掩码。最后，我们将根据您请求的输出强制转换为整数