Split CNN Feature Map

I have the Input image size: torch.Size([2, 3, 128, 128]), After 2nd Conv size is [2, 32, 128, 128] (nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=False) ) I want to divide this feature map into 4 feature block of [2,32,32,32], and pass to different attention Model. After combining it should be same as initial size [2, 32, 128, 128].

My question is, How to split the Feature map into 4 parts and add it back?

Thank you.

You could use F.fold/.unfold as seen e.g. here.

1 Like