No N-dimensional tranpose in PyTorch

Einops supports verbose transpositions for arbitrary number of dimensions:

from einops import rearrange
x  = torch.zeros(10, 3, 100, 100)
y  = rearrange(x, 'b c h w -> b h w c')
x2 = rearrange(y, 'b h w c -> b c h w') # inverse to the first

(and the same code works for tensorfow as well)


It's simply called differently in pytorch. torch.Tensor.permute will allow you to swap dimensions in pytorch like tf.transpose does in TensorFlow.

As an example of how you'd convert a 4D image tensor from NHWC to NCHW (not tested, so might contain bugs):

>>> img_nhwc = torch.randn(10, 480, 640, 3)
>>> img_nhwc.size()
torch.Size([10, 480, 640, 3])
>>> img_nchw = img_nhwc.permute(0, 3, 1, 2)
>>> img_nchw.size()
torch.Size([10, 3, 480, 640])

Tags:

Torch

Pytorch