PyTorch在对矩阵沿列方向求平均值时如何忽略掉padding的影响？

问题背景

这是最近某技术交流群群友提出来的一个问题，感觉挺有意思，因此做个简单的记录。

我们分两种情形来考虑，需要注意的是，无论是哪种情形，矩阵的形状基本是 (batch_size, sequence_length)。

情形一：整型矩阵

此时矩阵中的每个元素均为相应词元在词表中的索引，例如

a = torch.tensor([
    [3, 6, 1, 9, 9],
    [2, 7, 9, 9, 9],
])

其中填充词元在词表中的索引为 $9$。不难看出，第一个句子的实际长度为 $3$，第二个句子的实际长度为 $2$。

我们期望最终的输出结果为

tensor([3.3333, 4.5000])

即第一个句子的平均为 $(3+6+1)/3=3.333$，第二个句子的平均为 $(2+7)/2=4.5$。

情形二：浮点型矩阵

浮点型矩阵可能来源于神经网络中某一层的输出，例如

a = torch.tensor([
    [1.2607, -1.4688, 1.4340, 1.2454, 1.8901],
    [0.5616, -0.1035, 0.1797, 0.0235, -0.6699],
])

我们期望最终的输出结果为

tensor([0.4086, 0.2291])

即第一个句子的平均为 $(1.2607-1.4688+1.4340)/3=0.4086$，第二个句子的平均为 $(0.5616-0.1035)/2=0.2291$。

思路

我们先来看情形一。

为求得平均值，我们可以先统计每一行的实际长度

tensor([3, 2])

然后让填充词元所在位置都变成 $0$

tensor([[3, 6, 1, 0, 0],
        [2, 7, 0, 0, 0]])

对该矩阵沿列方向求和再除以上述的实际长度即可。

具体实现：

def mean_without_padding(a, padding_idx):
    b = (a != padding_idx).sum(dim=1)
    c = copy.deepcopy(a)
    c[a == padding_idx] = 0
    return c.sum(dim=1) / b

在计算浮点型矩阵的平均值时，我们一定会有情形一中的整型矩阵。这是因为，如果没有整型矩阵，我们就无法做embedding，进而得不到神经网络的输入，自然也就不可能产生浮点型矩阵。

大部分时候应该都是计算情形二中的矩阵，因为计算情形一并没有太多意义。