Pytorch maxpooling 通道维度

我试图用 Pytorch 构建一个 cnn，但在 maxpooling 方面遇到了困难。我考了斯坦福举办的cs231n。我记得，maxpooling可以用作维度推导步骤，例如，我有这个(1, 20, height, width)输入ot max_pool2d（假设我的batch_size是1）。如果我使用 (1, 1) 内核，我想得到这样的输出：(1, 1, height, width)，这意味着内核应该在通道维度上滑动。然而，在检查了 pytorch 文档后，它说内核在高度和宽度上滑动。感谢 Pytorch 论坛上的 @ImgPrcSng，他告诉我使用 max_pool3d，结果效果很好。但是 conv2d 层的输出和 max_pool3d 层的输入之间仍然存在重塑操作。所以很难聚合成 nn.Sequential，所以我想知道还有其他方法可以做到这一点吗？

像这样的东西会起作用吗？

from torch.nn import MaxPool1d
import torch.nn.functional as F


class ChannelPool(MaxPool1d):
    def forward(self, input):
        n, c, w, h = input.size()
        input = input.view(n, c, w * h).permute(0, 2, 1)
        pooled = F.max_pool1d(
            input,
            self.kernel_size,
            self.stride,
            self.padding,
            self.dilation,
            self.ceil_mode,
            self.return_indices,
        )
        _, _, c = pooled.size()
        pooled = pooled.permute(0, 2, 1)
        return pooled.view(n, c, w, h)

或者，使用einops

from torch.nn import MaxPool1d
import torch.nn.functional as F
from einops import rearrange


class ChannelPool(MaxPool1d):
    def forward(self, input):
        n, c, w, h = input.size()
        pool = lambda x: F.max_pool1d(
            x,
            self.kernel_size,
            self.stride,
            self.padding,
            self.dilation,
            self.ceil_mode,
            self.return_indices,
        )
        return rearrange(
            pool(rearrange(input, "n c w h -> n (w h) c")),
            "n (w h) c -> n c w h",
            n=n,
            w=w,
            h=h,
        )