[tensorflow笔记]-tensorflow实现带mask的reduce_mean

在使用tensorflow处理一些tensor时，有时需要对一个tensor取平均，可以使用tf.reduce_mean操作，但是这个没法处理带有mask的tensor数据，本文主要就是利用tensorflow的基本操作实现带mask的平均。

tf.reduce_mean

比如我们的数据是3维tensor，shape=（B，N，H），B表示batch_size、N表示最大长度、H表示向量维度，这样的3维tensor在NLP相关模型中很常见。

# input tensor
tensor = tf.constant([[[1, 2, 3], [3, 4, 5], [5, 6, 7]], 
                      [[1, 1, 2], [3, 3, 4], [5, 5, 6]]], dtype=tf.float32)
# reduce_mean
mean_out = tf.reduce_mean(tensor, axis=1)

print(mean_out)
'''
tf.Tensor(
[[3. 4. 5.]
 [3. 3. 4.]], shape=(2, 3), dtype=float32)
'''

带mask的reduce_mean

由于在NLP相关模型中，一个batch数据中，每个句子长度不同，有个mask参数表示句子长度，不够用padding补0操作，在取平均操作直接用tf.reduce_mean会把补0的位置也计算进去，下面实现的就是padding补0位置不参与平均值的计算。

# 下面两种写法都可以
def mask_reduce_mean(input_tensor, mask):
    mask = tf.dtypes.cast(mask, tf.float32)
    mask_sum = tf.reduce_sum(mask, axis=1, keepdims=True)
    mask_weight = tf.expand_dims(mask, axis=-1)
    output = tf.multiply(input_tensor, mask_weight)
    output = tf.reduce_sum(output, axis=1) / mask_sum
    return output

def mask_reduce_mean2(input_tensor, mask):
    mask = tf.dtypes.cast(mask, tf.float32)
    mask_sum = tf.reduce_sum(mask, axis=1, keepdims=True)
    mask_weight = tf.expand_dims(mask / mask_sum, axis=-1)
    output = tf.reduce_sum(input_tensor * mask_weight, axis=1)
    return output

测试数据

# input tensor
tensor = tf.constant([[[1, 2, 3], [3, 4, 5], [5, 6, 7]], 
                      [[1, 1, 2], [3, 3, 4], [5, 5, 6]]], dtype=tf.float32)
# input mask
# mask = np.array([[1, 1, 0], [1, 1, 1]])
mask = np.array([[True, True, False], [True, True, True]])

# mask reduce_mean
mean_out = mask_reduce_mean(tensor, mask)

print(mean_out)
'''
tf.Tensor(
[[2. 3. 4.]
 [3. 3. 4.]], shape=(2, 3), dtype=float32)
'''