class Pooling(nn.Module): def __init__(self): super().__init__() self.kernel = (2, 2) def forward(self, inputs): x = PoolFunc.apply(inputs.view(-1, *inputs.shape[2:]), self.kernel) return x.view(*inputs.shape[:2], *x.shape[1:])

这段代码实现了一个池化层的类，使用的是自定义的池化函数 PoolFunc。

在init函数中，定义了池化的 kernel 大小为 (2, 2)。

在 forward 函数中，首先将输入的 tensor 的 shape 调整为 (batch_size x channels x height x width)，然后调用 PoolFunc 进行池化操作，再将输出的 tensor 的 shape 调整回原本的形状。

需要注意的是，在这个类中并没有直接调用 PyTorch 自带的池化函数，而是通过自定义的 PoolFunc 进行操作，这是为了更加灵活地控制池化的方式。PoolFunc 可以是任何用户自己定义的函数，只需要保证输入和输出的 tensor 的 shape 是合法的即可。