The problem is:

I have a tensor, on this tensor i find the k_top elements, i create a mask with 0 if the element is not k_top and 1 otherwise, then i multiply the original tensor with the mask, creating tensor_masked. After that i calculate the average of the tensor.

The problem, now, is that i want to return a new tensor where only the elements of the k_top are equal to the average (-average if the original value inside the starting tensor was negative, average otherwise)

How can i do it?

This is the code for now:

def stc_compression(tensor, sparsification_rate):
    mask = tf.cast(tf.abs(tensor) >= tf.math.top_k(tf.abs(tensor), sparsification_rate)[0][-1], tf.float32)
    tensor_masked = tf.multiply(tensor, mask)
    average = tf.reduce_mean(tf.abs(tensor_masked)) / sparsification_rate
    return compressed_tensor

After that, is there a possible way to optimize this process?

CodePudding user response：

If you want to compute the average of the absolute values of the tensor_masked, use

tf.reduce_mean(tf.abs(tensor_masked))

if you want to compute the average of the absolute value of the top k values you should use

tf.reduce_sum(tf.abs(tensor_masked)) / sparsification_rate

Then to get the masked values with the right sign, you can use the tf.sign operation and use another mask to replace the values masked

def stc_compression(tensor, sparsification_rate):
        mask = tf.cast(tf.abs(tensor) >= tf.math.top_k(tf.abs(tensor), sparsification_rate)[0][-1], tf.float32)
        inv_mask = tf.cast(tf.abs(tensor) < tf.math.top_k(tf.abs(tensor), sparsification_rate)[0][-1], tf.float32)
        tensor_masked = tf.multiply(tensor, mask)
        average = tf.reduce_sum(tf.abs(tensor_masked)) / sparsification_rate
        compressed_tensor = tf.add( tf.multiply(average, mask) * tf.sign(tensor), tf.multiply(tensor, inv_mask))
        return compressed_tensor

Tell me if I misunderstood your problem.