Reshape channels into pixels with Pytorch

alagris 0
Viewed 18

By default pytorch encodes images as [batch_size, channels, height, width] tensors. I want to build a network that upsamples [b, c, h, w] image to [b, c, 2*h, 2*w]. The way I implement this is by first producing [b, 4*c, h, w] and then I want to reshape it to [b, c, 2*h, 2*w] in such a way that each 4*c-channel pixel becomes 4 neighbouring c-channel pixels such that the neighbours form a small square rather than a line of 4 pixels. Unfortunately the reshape() function arranges the neighbours into a line.

tensor reshape pytorch image
1 Answers

Consider using nn.Upsample: https://pytorch.org/docs/stable/generated/torch.nn.Upsample.html

You can use a scaling factor of 2 like in the doc. Not sure if the dimensions in this example are the same as yours, but inputs of the form N,C,H,W are also accepted.

>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
>>> input
tensor([[[[ 1.,  2.],
          [ 3.,  4.]]]])

>>> m = nn.Upsample(scale_factor=2, mode='nearest')
>>> m(input)
tensor([[[[ 1.,  1.,  2.,  2.],
          [ 1.,  1.,  2.,  2.],
          [ 3.,  3.,  4.,  4.],
          [ 3.,  3.,  4.,  4.]]]])
´´´
NiceGuySaysHi0
Related