像这样的工作吗？

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,
        )

要在所有通道上的每个坐标中最大化池，只需使用来自EINOP的layer

from einops.layers.torch import Reduce

max_pooling_layer = Reduce('b c h w -> b 1 h w', 'max')

层可以在您的模型中作为任何其他 Torch 模块使用

我不知道为什么其他的答案都这么复杂。在整个通道维度上的最大池化以获得仅具有1个通道声音的输出，相当于在该维度上取最大值：

torch.amax(left_images, dim=1, keepdim=True)