Published on

张量Tensor

Authors
  • Name
    Twitter

张量是存储算子输入数据与输出数据的容器。张量描述符（Tensor Desc）对输入与输出数据的描述，包含属性有：名称（Name），形状（Shape），数据类型（dtype），排布格式（Format）。

什么是张量？

张量（Tensor）是深度学习和机器学习中的核心数据结构，它是一个多维数组的数学抽象。在AI计算中，张量是存储算子输入数据与输出数据的容器，类似于传统编程中的数组，但具有更丰富的数学属性和操作能力。

张量的四个核心属性

1. 名称（Name）

  • 定义：张量的唯一标识符
  • 作用：用于在计算图中区分不同的张量，便于调试和可视化
  • 示例：input_tensor, weight_matrix, output_logits

2. 形状（Shape）

  • 定义：描述张量在每个维度上的大小
  • 表示方法：用元组表示，如 (batch_size, height, width, channels)
  • 常见形状：
    • 标量：() - 0维张量
    • 向量：(n,) - 1维张量
    • 矩阵：(m, n) - 2维张量
    • 3D张量：(batch, height, width) - 3维张量
    • 4D张量：(batch, height, width, channels) - 4维张量

3. 数据类型（dtype）

  • 定义：张量中元素的数据类型
  • 常见类型：
    • float32：32位浮点数（最常用）
    • float64：64位浮点数（双精度）
    • int32：32位整数
    • int64：64位整数
    • bool：布尔值
    • uint8：8位无符号整数

4. 排布格式（Format）

  • 定义：数据在内存中的存储和访问方式
  • 常见格式：
    • NCHW：(batch, channels, height, width) - PyTorch默认格式
    • NHWC：(batch, height, width, channels) - TensorFlow默认格式
    • CHW：(channels, height, width) - 无批次维度
    • HWC：(height, width, channels) - 无批次维度

张量的实际应用示例

图像处理中的张量

# 图像张量示例
# 形状：(batch_size, channels, height, width)
image_tensor = torch.randn(1, 3, 224, 224)  # 1张RGB图像，224x224像素
print(f"张量形状: {image_tensor.shape}")
print(f"数据类型: {image_tensor.dtype}")
print(f"设备: {image_tensor.device}")

自然语言处理中的张量

# 文本张量示例
# 形状：(batch_size, sequence_length, embedding_dim)
text_tensor = torch.randn(32, 128, 768)  # 32个句子，每个128个词，768维嵌入
print(f"张量形状: {text_tensor.shape}")

Tensor在内存中的存储方式

内存布局基础

张量在内存中的存储方式直接影响计算性能和内存使用效率。理解内存布局对于优化深度学习模型至关重要。

1. 连续内存存储

import torch

# 创建一个张量
tensor = torch.randn(2, 3, 4)
print(f"张量形状: {tensor.shape}")
print(f"是否连续: {tensor.is_contiguous()}")
print(f"步长(stride): {tensor.stride()}")
print(f"数据指针: {tensor.data_ptr()}")

2. 内存步长（Stride）

  • 定义：每个维度上相邻元素之间的内存距离
  • 计算：stride[i] = stride[i+1] * shape[i+1]
  • 示例：形状为(2, 3, 4)的张量，步长为(12, 4, 1)
# 步长示例
tensor = torch.randn(2, 3, 4)
print(f"形状: {tensor.shape}")
print(f"步长: {tensor.stride()}")

# 手动计算步长
manual_stride = (3*4, 4, 1)  # (12, 4, 1)
print(f"手动计算步长: {manual_stride}")

3. 内存对齐和数据类型

# 不同数据类型的内存占用
dtypes = [torch.float32, torch.float64, torch.int32, torch.int64, torch.uint8]
for dtype in dtypes:
    tensor = torch.tensor([1, 2, 3], dtype=dtype)
    print(f"{dtype}: 每个元素 {tensor.element_size()} 字节")

内存布局类型

1. 行主序（Row-Major）存储

  • 特点：最后一个维度在内存中连续存储
  • 优势：适合按行访问数据
  • 示例：C语言数组、PyTorch默认格式
# 行主序示例
tensor = torch.randn(2, 3)
print("行主序存储:")
for i in range(2):
    for j in range(3):
        print(f"tensor[{i},{j}] 在内存位置: {tensor[i,j].data_ptr()}")

2. 列主序（Column-Major）存储

  • 特点：第一个维度在内存中连续存储
  • 优势：适合按列访问数据
  • 示例：Fortran数组、某些科学计算库

内存优化策略

1. 内存连续性检查

# 检查张量是否连续
tensor = torch.randn(2, 3, 4)
print(f"原始张量连续: {tensor.is_contiguous()}")

# 转置操作可能破坏连续性
transposed = tensor.transpose(0, 2)
print(f"转置后连续: {transposed.is_contiguous()}")

# 恢复连续性
contiguous = transposed.contiguous()
print(f"恢复后连续: {contiguous.is_contiguous()}")

2. 内存预分配

# 预分配内存避免频繁分配
batch_size, seq_len, hidden_dim = 32, 128, 768

# 预分配张量
preallocated = torch.empty(batch_size, seq_len, hidden_dim, dtype=torch.float32)
print(f"预分配张量形状: {preallocated.shape}")
print(f"内存占用: {preallocated.numel() * preallocated.element_size()} 字节")

3. 内存池管理

# PyTorch内存池示例
import torch

# 创建大量小张量（会使用内存池）
tensors = []
for i in range(1000):
    tensor = torch.randn(10, 10)
    tensors.append(tensor)

# 查看内存使用情况
print(f"当前内存使用: {torch.cuda.memory_allocated() if torch.cuda.is_available() else 'CPU模式'}")

GPU内存存储

1. CUDA内存管理

# GPU内存存储示例
if torch.cuda.is_available():
    # 创建GPU张量
    gpu_tensor = torch.randn(1000, 1000).cuda()
    print(f"GPU张量设备: {gpu_tensor.device}")
    print(f"GPU内存使用: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
    
    # 内存清理
    del gpu_tensor
    torch.cuda.empty_cache()

2. 内存碎片化

# 内存碎片化示例
def create_fragmented_memory():
    tensors = []
    # 创建不同大小的张量
    for i in range(10):
        size = 1000 + i * 100
        tensor = torch.randn(size, size)
        tensors.append(tensor)
    
    # 删除部分张量
    del tensors[::2]  # 删除偶数索引的张量
    
    # 尝试分配大张量可能失败
    try:
        large_tensor = torch.randn(5000, 5000)
        print("大张量分配成功")
    except RuntimeError as e:
        print(f"大张量分配失败: {e}")

内存访问模式优化

1. 缓存友好的访问

# 缓存友好的内存访问
import time

def cache_friendly_access():
    tensor = torch.randn(1000, 1000)
    
    # 行优先访问（缓存友好）
    start_time = time.time()
    for i in range(1000):
        for j in range(1000):
            _ = tensor[i, j]
    row_time = time.time() - start_time
    
    # 列优先访问（缓存不友好）
    start_time = time.time()
    for j in range(1000):
        for i in range(1000):
            _ = tensor[i, j]
    col_time = time.time() - start_time
    
    print(f"行优先访问时间: {row_time:.4f}s")
    print(f"列优先访问时间: {col_time:.4f}s")
    print(f"性能差异: {col_time/row_time:.2f}x")

cache_friendly_access()

2. 内存带宽优化

# 内存带宽优化示例
def memory_bandwidth_test():
    sizes = [100, 500, 1000, 2000]
    
    for size in sizes:
        tensor = torch.randn(size, size)
        
        # 测试内存带宽
        start_time = time.time()
        result = tensor.sum()
        end_time = time.time()
        
        # 计算内存带宽
        memory_accessed = tensor.numel() * tensor.element_size() * 2  # 读取+写入
        bandwidth = memory_accessed / (end_time - start_time) / 1024**3  # GB/s
        
        print(f"大小 {size}x{size}: 带宽 {bandwidth:.2f} GB/s")

memory_bandwidth_test()

张量操作的重要性

  1. 内存效率：正确的排布格式可以显著提高内存访问效率
  2. 计算优化：硬件加速器（如GPU）对不同格式有优化偏好
  3. 框架兼容性：不同深度学习框架对张量格式有不同要求
  4. 调试便利：清晰的命名和形状信息有助于模型调试
  5. 内存管理：理解内存存储方式有助于优化内存使用和避免内存泄漏

总结

张量作为AI计算的基础数据结构，其四个核心属性（名称、形状、数据类型、排布格式）共同决定了数据在计算图中的存储、传输和处理方式。理解这些属性对于高效地进行深度学习模型开发和优化至关重要。

Discuss on Twitter • View on GitHub