No N-dimensional tranpose in PyTorch

25
votes

PyTorch's torch.transpose function only transposes 2D inputs. Documentation is here.

On the other hand, Tensorflow's tf.transpose function allows you to transpose a tensor of N arbitrary dimensions.

Can someone please explain why PyTorch does not/cannot have N-dimension transpose functionality? Is this due to the dynamic nature of the computation graph construction in PyTorch versus Tensorflow's Define-then-Run paradigm?

2
torchpytorch

2 Answers

40
votes

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])
3
votes

Einops supports verbose transpositions for arbitrary number of dimensions:

from einops import rearrange
x  = torch.zeros(10, 100, 100, 3)
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)