pytorch交换tensor的指定维度

pytorch中有两种方式可以实现tensor指定维度的交换,第一个是torch.permute(),第二个方法是torch.transpose()。

二者不同是torch.permute()可以同时交换多个维度,而torch.transpose()每次只能交换两个维度。

方式一:torch.permute()

参数列表:

  • input:待交换的张量
  • dims:需要交换维度的索引

该函数会按照我们指定维度方式重新排列,例如我们下面定义了一个张量维度为【2,3,4】,如果我们要将维度变为【4,3,2】,就需要交换第一个维度和第三个维度,那我们传入的参数维度索引就应该为【2,1,0】,该索引对应维度的顺序,原来是【0,1,2】,现在是【2,1,0】。

a = torch.randn(2, 3, 4)
print(a.shape)

print(torch.permute(a, (2, 1, 0)).shape)
torch.Size([2, 3, 4])
torch.Size([4, 3, 2])

或者还可以使用transpose函数直接交换维度

方式二:torch.transpose()

对于该方法每次只能交换两个维度,输入的参数很简单就是需要交换的两个维度的索引。

print(torch.transpose(a, 0, 1).shape)
torch.Size([3, 2, 4])

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
xiaoxingxing的头像xiaoxingxing管理团队
上一篇 2023年8月8日
下一篇 2023年8月8日

相关推荐