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

我有一个逻辑问题

我正在尝试创建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)

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

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

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

---

显然，除了PyTorch张量外，这个问题还可以应用于张量流张量或NumPy

ndarray

s。我和这三个人一起工作，逻辑大概是一样的。

？@hkchengrex-Oh，很有魅力！天哪，我真不敢相信我花了两个小时来想办法。我不擅长谷歌搜索。

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