按所需输出形状划分的平均池

Question

我有一个逻辑问题

我试图创建PyTorch的nn.functional.avg_pool函数的某种版本，只做了一点修改：我希望根据输出的所需大小动态选择填充/跨步。我已经成功地为单个1d张量计算出了这一点：

def ave_pool(tnsr, reduction, pad_val=0, dim = 1):

    size = len(tnsr)
    if size % reduction != 0:
        padlen = np.ceil(size / reduction) * reduction - size
        fpad, bpad = np.ceil(padlen / 2), np.floor(padlen/2)
        tnsr = F.pad(tnsr, (int(fpad), int(bpad)), "constant", pad_val)
        size += padlen
    group = int(size / reduction)
    return tnsr.reshape((-1, group)).mean(dim = 1)

这对于单个一维张量来说非常有效。然而，当一维张量通过我的神经网络时（此时基本上是二维输入），我想把这个想法应用到“批量”的一维张量上

# This works as expected:
x = torch.tensor([2,3,4,8,9,10]).float()
ave_pool(x, 3)
>>> 
tensor([2.5000, 6.0000, 9.5000])

# But what's the logic to up the dimensions to two?
x = torch.tensor([[2,3,4,8,9,10],
                 [6,7,8,9,19,12]]).float()
ave_pool(x, 3)
>>> 
tensor([ 0.,  2.,  3.,  4.,  8.,  9., 10.,  0.,  6.,  7.,  8.,  9., 19., 12.])

那么，我如何将当前工作代码的逻辑扩展到2d张量，以便前一个代码块的输出为：

tensor([[2.5000, 6.0000, 9.5000],
        [ 6.5000,  8.5000, 15.5000]])

显然，除了PyTorch张量外，这个问题还可以应用于张量流张量或NumPyndarray。我和这三个人一起工作，逻辑大概是一样的