| import torch | |
| import torch.nn as nn | |
| class ChannelCompression(nn.Module): | |
| """ | |
| Reduces the input to 2 channels by concatenating the global average pooling and global max pooling outputs. | |
| In: HxWxC | |
| Out: HxWx2 | |
| """ | |
| def forward(self, x): | |
| return torch.cat( | |
| (torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1 | |
| ) | |