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

【问题标题】：More efficient way to create a mask of top `k` elements in tensorflow在 tensorflow 中创建顶部“k”元素掩码的更有效方法
【发布时间】：2021-04-20 15:35:39
【问题描述】：

我想创建一个函数，对于给定的1d-张量输出掩码，在与顶部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'

【问题讨论】：

标签： python arrays tensorflow

【解决方案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]。然后我们通过询问>= 问题来创建掩码。最后，我们根据您要求的输出转换为整数

【讨论】：