How to reshape tensor in a right way?

1

How to reshape torch.Size([2, 3, 4, 5]) to torch.Size([2, 5, 3, 4]) and then get back to torch.Size([2, 3, 4, 5])?
In my case torch.Size([batch_size, x_dim, y_dim, 3_RGB]). I need convert to torch.Size([batch_size, 3_RGB, x_dim, y_dim]) and use it with CNN layer and then get back to original size: torch.Size([batch_size, x_dim, y_dim, 3_RGB])

2

For example,

x = torch.rand(2,3,4,5) # now x.shape is torch.Size([2, 3, 4, 5])
y = x.permute(0,3,1,2) # now y.shape is torch.Size([2, 5, 3, 4])
    # which means: take the x's dimension 0,3,1,2 to be y's dimension 0,1,2,3
    # i.e. x's dimension 0 is 2; dimension 1 is 3; dimension 2 is 4; dimension 3 is 5.
    # so x's dimension 0,3,1,2 would be 2,5,3,4.
z = y.permute(0,2,3,1) # now z.shape is torch.Size([2, 4, 3, 5])
1 Like