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

谁借莪１个温暖的怀抱￠ 2022-11-01 10:56 140阅读 0赞

在使用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) '''