PyTorch中的movedim、transpose与permute

news2025/2/2 22:47:17

在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])