理解pytorch系列：transpose是怎么实现的

在PyTorch中，

transpose()

是一种操作，它交换张量中两个指定维度的位置。实现这一点的关键在于不实际移动数据，而是通过改变张量的元数据（包括步长（stride）和尺寸（size））来达到效果。

举例来说，假设我们有一个形状为

(3, 4)

的二维张量，其内存布局为行优先（row-major）即C风格的。当我们对这个张量执行

transpose(0, 1)

操作时，我们期望该张量行变成列，列变成行，即得到一个形状为

(4,3)

的新视图。

这是通过以下步骤完成的：

1. 改变尺寸：改变size元数据，使得原本第一个维度（行）的大小与第二个维度（列）的大小交换。
2. 改变步长：步长（stride）是一个数组，指示了在每个维度上移动一个元素需要跳过的内存位置数。执行transpose()时，交换了两个维度的步长。在行优先存储的张量中，行的步长通常比列的步长大。
3. 不移动数据：实际上数据并没有在内存中移动，只是改变了在这块内存空间上的解释方式。

以下是一个简单的示例：

import torch

# 创建一个 3x4 的张量
x = torch.arange(12).view(3,4)print("Original tensor:")print(x)# 输出：# tensor([[ 0,  1,  2,  3],#         [ 4,  5,  6,  7],#         [ 8,  9, 10, 11]])# 现在使用 transpose 来交换两个维度
y = x.transpose(0,1)print("\nTransposed tensor:")print(y)# 输出：# tensor([[ 0,  4,  8],#         [ 1,  5,  9],#         [ 2,  6, 10],#         [ 3,  7, 11]])

在这个例子中，

x

的形状发生了变化，但它的内存布局没有改变。通过调整步长和大小，

transpose()

创建了一个新的张量视图。

在PyTorch的底层C++实现中，同样接口会调用ATen库（张量操作库，是PyTorch的核心）中的对应函数，ATen函数会修改张量对象所关联的元数据以实现

transpose()

操作。这意味着实际的CPU或GPU中的数据不会因为

transpose()

操作而移动。这种"懒惰"操作提高了性能，特别是对于大型张量，因为它避免了不必要的数据拷贝。