More efficient way to create a mask of top `k` elements in tensorflow

user_13614416 0
Viewed 403

I would like to create a function, that for a given 1d-tensor outputs the mask, where on the places, corresponding to the top k values there are 1 and 0 elsewhere. Namely, I have for example:

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

I've created the following function:

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

However for the case of interest it works rather slowly, and as I've measured most of the time is dominated by tf.tensor_scatter_nd_update. What would be the faster alternative?

The typical size of tensor is 10^3-10^4 elements and k is of order `10^2-10^3'

tensorflow arrays python
1 Answers

I'd go for finding the top K value, following by a size comparison

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)>

Explanation

tf.math.top_k returns two values, the first is a tensor with the actual top k values, and the second are the indices. We take the values, and then access [-1] which is the smallest value. Then we create the mask by asking the >= question. Finally we cast to integer per your requested output

user_53680830
Related