张量实现细节

创建和操作PyTorch张量通常感觉很直观，但了解其内部结构有助于编写高效代码、调试复杂行为以及构建自定义操作。torch.Tensor不仅仅是一个多维数组；它是一个精密的D对象，包含定义内存中原始数值数据如何被解释的元数据。

张量对象与内存存储

本质上，每个PyTorch张量都持有一个指向torch.Storage对象的引用。可以将torch.Storage视为一个连续的、一维的特定类型（例如，float32、int64）数值数据数组。Tensor对象本身不直接包含数值，而是持有描述如何在与其相关联的Storage中查看数据的元数据。

这种分离很重要，因为多个张量可以共享相同的内存存储。切片、转置或重塑等操作通常会创建具有不同元数据的新Tensor对象，但它们指向由Storage管理的相同内存块。这使得这些操作非常节省内存，因为它们通常不涉及数据复制。

import torch

# 创建一个张量
x = torch.arange(12, dtype=torch.float32)
print(f"原始张量 x: {x}")

# 存储是一个包含12个浮点数的一维数组
print(f"存储元素: {x.storage().tolist()}")
print(f"存储类型: {x.storage().dtype}")
print(f"存储大小: {len(x.storage())}")

# 通过重塑创建视图
y = x.view(3, 4)
print(f"\n重塑后的张量 y:\n{y}")

# y 具有不同的形状/步幅，但共享相同的存储
print(f"y 是否与 x 共享存储？ {y.storage().data_ptr() == x.storage().data_ptr()}")

# 修改视图会影响原始张量（反之亦然）
y[0, 0] = 99.0
print(f"\n修改后的 y:\n{y}")
print(f"修改 y 后的原始 x: {x}")

在上面的例子中，x和y是不同的Tensor对象，但由于y是通过reshape创建的视图，它们共享相同的内存存储。修改y中的元素也会改变通过x可见的相应元素。

张量元数据

除了对其Storage的引用之外，Tensor对象还维护多项元数据，这些数据定义了其属性和数据解释方式：

1. 设备（device）： 指定张量数据所在的设备，可以是CPU（torch.device('cpu')）或特定的GPU（torch.device('cuda:0')）。张量间的操作通常需要数据在相同设备上。在设备之间移动数据（例如，使用.to(device)）涉及内存复制，这可能影响性能。

2. 数据类型（dtype）： 定义张量中元素的数值类型，例如torch.float32、torch.int64、torch.bool。操作通常要求张量具有兼容的数据类型，并且数据类型的选择显著影响内存使用和数值精度。

3. 形状（shape 或 size()）： 一个表示张量维度的元组。例如，一个3x4矩阵的形状是(3, 4)。

4. 存储偏移（storage_offset()）： 一个整数，表示此张量数据在内存存储中开始的索引。对于直接创建的张量（非视图），这通常是0。切片可能具有非零偏移量。

5. 步幅（stride()）： 这也许是理解内存布局最重要的元数据。步幅是一个元组，其中第 $i$ 个元素指定了沿张量第 $i$ 维度移动一步所需在内存（Storage中的元素数量）中的步长。

考虑一个3x4的张量t：

t = torch.arange(12, dtype=torch.float32).view(3, 4)
print(f"张量 t:\n{t}")
print(f"形状: {t.shape}")
print(f"步幅: {t.stride()}")

输出：

张量 t:
tensor([[ 0.,  1.,  2.,  3.],
        [ 4.,  5.,  6.,  7.],
        [ 8.,  9., 10., 11.]])
形状: torch.Size([3, 4])
步幅: (4, 1)

步幅(4, 1)表示：

• 沿维度0（向下行）移动一步，你需要在内存一维Storage中跳跃4个元素。（例如，从元素0到元素4）。
• 沿维度1（横跨列）移动一步，你需要在内存一维Storage中跳跃1个元素。（例如，从元素0到元素1）。

步幅确定了多维张量如何映射到线性Storage。

内存布局：连续与非连续

如果张量的元素在Storage中的排列顺序与标准的C风格（行主序）遍历顺序相同，则该张量在内存中被认为是连续的。对于连续张量，步幅通常遵循一种模式：最后一个维度的步幅是1，倒数第二个维度的步幅是最后一个维度的大小，以此类推。对于我们上面3x4的张量t，步幅是(4, 1)，这符合这种模式（stride[1] == 1，stride[0] == shape[1] == 4），因此它是连续的。

print(f"t 是否连续？ {t.is_contiguous()}") # 输出: True

然而，转置等操作可以创建非连续张量（视图）。

t_transposed = t.t() # 转置操作
print(f"\n转置后的张量 t_transposed:\n{t_transposed}")
print(f"形状: {t_transposed.shape}")
print(f"步幅: {t_transposed.stride()}")
print(f"t_transposed 是否连续？ {t_transposed.is_contiguous()}") # 输出: False
print(f"t_transposed 是否与 t 共享存储？ {t_transposed.storage().data_ptr() == t.storage().data_ptr()}") # 输出: True

输出：

转置后的张量 t_transposed:
tensor([[ 0.,  4.,  8.],
        [ 1.,  5.,  9.],
        [ 2.,  6., 10.],
        [ 3.,  7., 11.]])
形状: torch.Size([4, 3])
步幅: (1, 4)
t_transposed 是否连续？ False
t_transposed 是否与 t 共享存储？ True

请注意，t_transposed的形状是(4, 3)，但其步幅是(1, 4)。沿维度0（在转置视图中向下行）移动，在原始存储中跳跃1个元素。沿维度1（横跨列）移动，跳跃4个元素。这种布局不是C语言连续的。

为什么连续性很重要？

• 性能： 许多PyTorch操作（特别是低层CPU/GPU核）针对连续张量进行了优化。当数据在内存中按顺序排列时，它们可以更高效地处理数据。非连续张量可能触发内部内存复制或使用较慢的算法。
• 兼容性： 某些操作，如view()，要求张量是连续的。如果你尝试在非连续张量（如t_transposed）上使用view()，会得到错误。在这种情况下，你通常需要使用reshape()，如果可能它会返回一个视图，但如果需要满足形状改变，它会返回一个副本。或者，你可以使用.contiguous()方法显式创建一个连续的副本。

# 这会引发 RuntimeError，因为 t_transposed 不连续
# flat_view = t_transposed.view(-1)

# .contiguous() 在需要时会创建一个具有连续内存布局的新张量
t_contiguous_copy = t_transposed.contiguous()
print(f"\n连续副本是否连续？ {t_contiguous_copy.is_contiguous()}") # 输出: True
print(f"连续副本的步幅: {t_contiguous_copy.stride()}") # 输出: (3, 1)
print(f"存储共享？ {t_contiguous_copy.storage().data_ptr() == t_transposed.storage().data_ptr()}") # 输出: False (这是一个副本)

# 现在视图可以工作了
flat_view = t_contiguous_copy.view(-1)
print(f"展平视图: {flat_view}")

以下图表说明了两个张量，T（原始3x4）和T_transpose（其转置），如何将其元素映射到相同的底层一维存储块。请注意步幅如何决定不同的访问模式。

张量元数据与内存存储的关系，适用于3x4张量T及其转置T_transpose。两个张量对象都指向相同的存储，但根据它们的形状、步幅和偏移量以不同方式解释它。箭头表示元素如何从张量视图映射到存储索引。

理解这些实现细节，Tensor元数据与Storage的区别，步幅的作用，以及连续性的思想，为推断PyTorch中内存使用、性能特点和各种张量操作行为提供了坚实的支持。当优化瓶颈或与低层代码交互时，这些知识尤其有用。