Pooling over channels in pytorch

Question

In tensorflow, I can pool over the depth dimension which would reduce the channels and leave the spatial dimensions unchanged. I'm trying to do the same in pytorch but the documentation seems to say pooling can only be done over the height and width dimensions. Is there a way I can pool over channels in pytorch? I've a tensor of shape [1,512,50,50] I'm trying to use pooling to bring the number of channels down to 3. I saw this question but did not find the answer helpful.

Answer 1

The easiest way to reduce the number of channels is using a 1x1 kernel:

import torch                                                                                                                                                                                               

x = torch.rand(1, 512, 50, 50)                                                                                                                                                                                  
conv = torch.nn.Conv2d(512, 3, 1)                                                                                                                                                                             
y = conv(x)                                                                                                                                                                                                   

print(y.size())                                                                                                                                                                                                   
# torch.Size([1, 3, 50, 50])

If you really need to perform pooling along the channels dimension due to some reason, you may want to permute the dimensions so that the channels dimension is swapped with some other dimension (e.g. width). This idea was referenced here.