【python深度学习】——tensor内部存储结构|内存优化与as_strided|内存紧凑化contiguous
- 1. tensor的元数据(metadata)和存储区(storage)
- 1.1 元数据(metadata)
- 1.2 存储区(Storage)
- 1.3 查看tensor信息的代码示例
- 2. 通过步长(Stride)、偏移量(storage_offset)以及“as_strided”实现内存优化
- 2.1 步长(Stride)、偏移量(storage_offset)避免不必要的数据拷贝原理
- 2.2 示例代码——通过as_strided来实现高级视图操作
- 3. 内存紧凑化方法contiguous
1. tensor的元数据(metadata)和存储区(storage)
pytorch中的tensor内部存储结构分为两个部分——元数据(metadata)和存储区(storage)
1.1 元数据(metadata)
元数据中包含以下信息:
- 形状——size(也是shape)
- 数据类型——dtype
- 设备——device
- 步长——stride。步长告诉我们, 在内存中,移动到该tensor的下一个元素需要移动多少个位置。
- 内存偏移—— storage_offset。当前 Tensor 的第一个元素在其底层存储区中的起始位置, storage_offset 的概念对于理解视图操作(如切片、reshape 等)非常重要。
通过调整 storage_offset 和步长(stride),我们可以高效地重新排列和访问数据,而不需要实际的数据拷贝, 这些在后面来详细介绍
1.2 存储区(Storage)
存储区是实际存储 Tensor 数据的地方。它包含了 Tensor 中所有元素的实际值,是一个一维的连续内存区, 它是可以共享的!
存储区有以下几点特性:
- 连续性:虽然 Tensor 可以是多维的,但存储区是一维的,所有元素按照特定顺序存储在内存中。步长(stride)和形状(shape)帮助我们理解这些元素在多维空间中的排列方式。
- 引用计数:存储区可以被多个 Tensor 共享。引用计数机制用于跟踪有多少个 Tensor 引用了同一个存储区,这有助于管理内存。
- 数据类型一致:存储区中的所有数据元素类型一致,并且由 dtype 指定。
1.3 查看tensor信息的代码示例
我们可以这样查看tensor的信息:
import torch
# 创建一个形状为 (3, 4) 的 Tensor
tensor = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=torch.float32)
# 查看 Tensor 的元数据
print(f"形状: {tensor.shape}") #等价于print(f"形状: {tensor.size}")
print(f"数据类型: {tensor.dtype}")
print(f"设备: {tensor.device}")
print(f"步长: {tensor.stride()}")
print(f"存储区: {tensor.storage_offset()}")
# 查看存储区信息
print(f"存储区: {tensor.storage()}")
print(f"存储区数据类型: {tensor.storage().dtype}")
print(f"存储区元素数: {tensor.storage().size()}")
2. 通过步长(Stride)、偏移量(storage_offset)以及“as_strided”实现内存优化
2.1 步长(Stride)、偏移量(storage_offset)避免不必要的数据拷贝原理
深度学习中数据量往往是非常庞大的, 因此, tensor的storage是可以共享的, 如下图所示——tensor A和B的元素是相同的, 在内存中,其实是指向同一片空间的, 只是通过stride的不同,来实现呈现的tensor不同。
2.2 示例代码——通过as_strided来实现高级视图操作
# 创建一个 3x4 的 Tensor
tensor = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=torch.float32)
# 使用 as_strided 创建一个新的视图
as_strided_tensor = tensor.as_strided((2, 2), (4, 2), storage_offset=0)
print(f"as_strided Tensor:\n{as_strided_tensor}") # output: tensor([[1,3],[5,7]])
as_strided_tensor = tensor.as_strided((2, 2), (4, 2), storage_offset=1)
print(f"as_strided Tensor:\n{as_strided_tensor}") # output: tensor([[2,4],[6,8]])
3. 内存紧凑化方法contiguous
由前面两节我们已经知道, 有些情况下, tensor数据只是视图不同, 元素的存储其实没有改变, 那么在反复操作时, 有可能带来计算效率的下降。
那么当我们完成了视图操作并希望得到一个新的紧凑的 Tensor,可以使用 contiguous 方法将视图变为内存连续的, 例如下面的代码:
# 创建一个转置视图
transposed_tensor = tensor.t()
# 将转置视图变为连续的 Tensor
contiguous_tensor = transposed_tensor.contiguous()
print(f"连续 Tensor:\n{contiguous_tensor}")
print(f"连续 Tensor 的步长: {contiguous_tensor.stride()}")
在这个例子中,contiguous_tensor 是一个新的 Tensor,数据被复制到一个新的连续存储区中。