PyTorch学习之torch.transpose函数

PyTorch学习之torch.transpose函数

一、简介

torch.transpose 函数我们用于交换张量的维度。

二、语法

torch.transpose 函数用于交换给定张量的两个维度，其语法如下：

torch.transpose(input, dim0, dim1)

三、参数

• input：待交换维度的张量。
• dim0：第一个要交换的维度。
• dim1：第二个要交换的维度。

四、示例

示例 1：基本用法

import torch

# 创建一个 2x3 的张量
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
print("原始张量：")
print(x)

# 交换第 0 维和第 1 维
# x_transposed = torch.transpose(x, 0, 1)
x_transposed = x.transpose(0, 1)
print("交换维度后的张量：")
print(x_transposed)

输出：

原始张量：
tensor([[1, 2, 3],
        [4, 5, 6]])
交换维度后的张量：
tensor([[1, 4],
        [2, 5],
        [3, 6]])

在这个例子中，原始的 2x3 张量通过交换维度变成了 3x2 张量。

示例 2：高维张量的维度交换

# 创建一个 2x3x4 的张量
x = torch.arange(24).reshape(2, 3, 4)
print("原始张量：")
print(x)

# 交换第 0 维和第 2 维
x_transposed = torch.transpose(x, 0, 2)
print("交换维度后的张量：")
print(x_transposed)

输出：

原始张量：
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])
交换维度后的张量：
tensor([[[ 0, 12],
         [ 4, 16],
         [ 8, 20]],

        [[ 1, 13],
         [ 5, 17],
         [ 9, 21]],

        [[ 2, 14],
         [ 6, 18],
         [10, 22]],

        [[ 3, 15],
         [ 7, 19],
         [11, 23]]])

在这个例子中，原始的 2x3x4 张量通过交换第 0 维和第 2 维变成了 4x3x2 张量。