PyTorch中的movedim、transpose与permute

在PyTorch中，movedim、transpose 和 permute这三个操作都可以用来重新排列张量（tensor）的维度，它们功能相似却又有所不同。

movedim

🔗 torch.movedim

• 用途：将张量的一个或多个维度移动到新的位置。
• 参数：需要指定要移动的维度和目标位置。
• 特点：
  • 可以移动单个或多个维度。
  • 可以指定移动后的维度顺序。
  • 移动维度时，不会改变其他维度的相对顺序。

movedim 的重点是移动维度到指定位置，而transpose的重点是交换两个维度，这是它们之间最显著的区别，此外movedim还可以移动多个维度。

>>> x = torch.randint(1,100,(1,2,3,4))
>>> x
tensor([[[[58,  1, 62, 12],[65, 99, 20, 15],[95,  6, 54, 67]],[[92, 62, 97, 15],[17, 19, 70, 48],[17, 90, 76, 96]]]])
>>> x.shape
torch.Size([1, 2, 3, 4])
>>> ## 移动单个维度
>>> x.movedim(0,3).shape    # 把原tensor中的dim0移动到dim3
torch.Size([2, 3, 4, 1])
>>> x.transpose(0,3).shape  # 交换dim-0 和 dim-3
torch.Size([4, 2, 3, 1])
>>> 
>>> ## 同时移动多个维度
>>>> x.movedim((0,2),(1,3)).shape # 把原tensor中的dim0移动到dim1， dim2移动到dim3
torch.Size([2, 1, 4, 3])
>>> x.movedim((0,1),(1,2)).shape  # 把原tensor中的dim0移动到dim1， dim1移动到dim2
torch.Size([3, 1, 2, 4])

总结

• movedim是把指定维度移动到新的维度位置，其他未移动的的维度保持相对位置不变。
• 对于 source_tensor.movedim(d1, d2):
  • 如果 d1,d2 是连续的，则效果等同于 source_tensor.transpos(d1, d2)
  • 如果 d1 < d2，则移动后，原始维度中（d1, d2] 范围内的维度左移一位；
  • 如果 d1 > d2，则移动后，原始维度中（d1, d2] 范围内的维度右移一位；

transpose

🔗 torch.transpose

• 用途：交换张量的两个维度。
• 参数：需要指定要交换的两个维度的索引。
• 特点：
  • 只能交换两个维度。
  • 对于2D张量，相当于矩阵的转置。
  • 对于高维张量，只能进行一次交换。

>>> x = torch.randint(2,300,(2,3,4))
>>> x
tensor([[[ 23, 266, 112, 126],[176, 290, 291, 134],[155, 214, 238, 208]],[[ 20, 270, 192, 254],[ 59, 120,  47,   5],[247, 173,  94,  28]]])
>>> x.transpose(0,2).shape
torch.Size([4, 3, 2])
>>> x.transpose(0,2)
tensor([[[ 23,  20],[176,  59],[155, 247]],[[266, 270],[290, 120],[214, 173]],[[112, 192],[291,  47],[238,  94]],[[126, 254],[134,   5],[208,  28]]])

总结

• torch.transpose 是用于交换两个维度，且只能交换两个维度。
• torch.permute 可以交换多个维度

permute

🔗torch.permute

• 用途：重新排列张量的所有维度。
• 参数：需要提供一个维度索引的排列。
• 特点：
  • 可以重新排列所有维度的顺序。
  • 参数是一个包含所有维度索引的元组，且每个索引只能出现一次。

>>> x = torch.randint(2,300,(2,3,4))
>>> x
tensor([[[158, 169,  28,  85],[212,  14, 265, 119],[  6, 201, 221, 145]],[[114, 209, 272, 232],[ 46, 179, 121, 296],[ 89,  76, 142, 185]]])
>>> x.shape
torch.Size([2, 3, 4])
>>> x.permute(2,0,1).shape  # 表示：原tensor的dim2放到输出tensor的dim0,dim0放到dim1,dim1放到dim2。
torch.Size([4, 2, 3])

异同点

相同点：

• 都用于改变张量的维度顺序。
• 都不会改变张量中的数据，只是改变了数据的布局。

不同点：

• transpose仅限于交换两个维度，而movedim和permute可以处理多个维度。
• permute需要指定所有维度的排列顺序，而movedim只需要指定要移动的维度和目标位置。
• movedim在移动维度时，其他维度的相对顺序保持不变，而permute会根据提供的排列顺序完全重新排列所有维度。