Tensorflow从一个张量和一个掩码中创建一个新的张量

Question

问题是:

我有一个张量，在这个张量上我找到k_top元素，如果元素不是k_top，我创建一个0的掩码，否则为1，然后我将原始张量与掩码相乘，创建tensor_mask。然后计算张量的平均值。

现在的问题是，我想返回一个新的张量，其中只有k_top的元素等于平均值(如果初始张量内部的原始值为负，则为-平均值，否则为+平均值)

我该怎么做呢?

这是现在的代码:

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

在那之后，是否有可能优化这个过程?

Answer 1

如果要计算tensor_mask的绝对值的平均值，请使用

tf.reduce_mean(tf.abs(tensor_masked))

如果你想计算前k个值的绝对值的平均值，你应该使用

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

然后，要获得具有正确符号的掩码值，可以使用tf。符号操作并使用另一个掩码来替换掩码

的值

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

如果我误解了你的问题，请告诉我。