Pytorch 原理 – 张量存储及访问

本文重点探讨,Pytorch 中张量在内存中的存储布局以及对应的若干变换操作 。

张量(Tensor)是数学和物理学中用于表示多维数组的一个概念,在机器学习、深度学习等领域也得到了广泛应用,简单来说,张量可以被视为一种广义的矩阵,它可以拥有任意数量的维度。我们生成一个3*3 的张量, 代码段1。

import torch
t = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
print(t, t.shape)

output:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]) 
        
torch.Size([3, 3]  # 3*3

代码段1,展示了一个二维张量,3行3列。我们知道,计算机的内存(RAM)本质上是一个线性、一维的存储空间。每一个存储单元(通常是字节,byte)都有一个唯一的地址,这些地址从 0 开始顺序排列,形成一个连续的一维空间。因此不论张量是多少维的,最终都会被映射为一维。在映射机制里,涉及两个重要概念,连续性(contiguity rule)步长(strides)

连续性:数据元素在内存空间的排列顺序,主要包括行优先和列优先两类,Pytorch 默认使用行优先方式。

行优先(以C语言为例)

在行优先顺序(Row-major order)中,,内存先存储第 0 行的所有元素,接着是第 1 行的所有元素,依此类推。

对于代码1中的矩阵,其在内存中行优先布局:

1,2,3,4,5,6,7,8,9

列优先(以Fortran为例)

在列优先顺序(Column-major order)中,内存先存储第 0 列的所有元素,接着是第 1 列的所有元素,依此类推。

对于代码1中的矩阵,其在内存中列优先布局:

1,4,7,2,5,8,3,6,9

数据在内存中的存储布局决定了读取效率,代码段2展示了c++ 中二维张量行优先和列优先的访问差异。

#include <iostream>
#include <chrono>

// loop1 accesses data in matrix 'a' in row major order,
// since i is the outer loop variable, and j is the
// inner loop variable.
int loop1(int a[4000][4000]) {
 int s = 0;
 for (int i = 0; i < 4000; ++i) {
   for (int j = 0; j < 4000; ++j) {
     s += a[i][j];
   }
 }
 return s;
}

// loop2 accesses data in matrix 'a' in column major order
// since j is the outer loop variable, and i is the
// inner loop variable.
int loop2(int a[4000][4000]) {
 int s = 0;
 for (int j = 0; j < 4000; ++j) {
   for (int i = 0; i < 4000; ++i) {
     s += a[i][j];
   }
 }
 return s;
}

int main() {
 static int a[4000][4000] = {0};
 for (int i = 0; i < 100; ++i) {
   int x = rand() % 4000;
   int y = rand() % 4000;
   a[x][y] = rand() % 1000;
 }

 auto start = std::chrono::high_resolution_clock::now();
 auto end = start;
 int s = 0;

#if defined RUN_LOOP1
 start = std::chrono::high_resolution_clock::now();

 s = 0;
 for (int i = 0; i < 10; ++i) {
   s += loop1(a);
   s = s % 100;
 }
 end = std::chrono::high_resolution_clock::now();

 std::cout << "s = " << s << std::endl;
 std::cout << "Time for loop1: "
   << std::chrono::duration<double, std::milli>(end - start).count()
   << "ms" << std::endl;
#endif

#if defined RUN_LOOP2
 start = std::chrono::high_resolution_clock::now();
 s = 0;
 for (int i = 0; i < 10; ++i) {
   s += loop2(a);
   s = s % 100;
 }
 end = std::chrono::high_resolution_clock::now();

 std::cout << "s = " << s << std::endl;
 std::cout << "Time for loop2: "
   << std::chrono::duration<double, std::milli>(end - start).count()
   << "ms" << std::endl;
#endif
}

编译并执行代码段2,可以看到行优先方式相比较于列优先方式,执行效率提高了近6倍。

g++ -O2 a.cpp -DRUN_LOOP1 -DRUN_LOOP2
./a.out 
s = 70
Time for loop1: 26.414ms
s = 70
Time for loop2: 143.802ms

效率的提高无非是指令的执行和数据的读取,我们可以使用Cachegrind工具针对缓存行为进行分析。(Cachegrind 主要专注于模拟程序执行期间的缓存行为)。

代码段3 针对loop1 中(行优先)的指令和数据缓存行为分析,指令缓存命中率和数据缓存命中率都很高。

g++ -O2 a.cpp -DRUN_LOOP1
valgrind --tool=cachegrind ./a.out

==1701994== Cachegrind, a cache and branch-prediction profiler
==1701994== Copyright (C) 2002-2017, and GNU GPL'd, by Nicholas Nethercote et al.
==1701994== Using Valgrind-3.18.1 and LibVEX; rerun with -h for copyright info
==1701994== Command: ./a.out
==1701994== 
--1701994-- warning: L3 cache found, using its data for the LL simulation.
^@s = 70
Time for loop1: 458.947ms
==1701994== 
==1701994== I   refs:      162,960,511
==1701994== I1  misses:          2,237
==1701994== LLi misses:          2,236
==1701994== I1  miss rate:        0.00%
==1701994== LLi miss rate:        0.00%
==1701994== 
==1701994== D   refs:       40,771,337  (40,567,464 rd   + 203,873 wr)
==1701994== D1  misses:     10,016,418  (10,013,794 rd   +   2,624 wr)
==1701994== LLd misses:     10,009,854  (10,008,097 rd   +   1,757 wr)
==1701994== D1  miss rate:        24.6% (      24.7%     +     1.3%  )
==1701994== LLd miss rate:        24.6% (      24.7%     +     0.9%  )
==1701994== 
==1701994== LL refs:        10,018,655  (10,016,031 rd   +   2,624 wr)
==1701994== LL misses:      10,012,090  (10,010,333 rd   +   1,757 wr)
==1701994== LL miss rate:          4.9% (       4.9%     +     0.9%  )

代码段4 针对loop2中(列优先)的指令和数据缓存行为分析,指令缓存命中率很高,但是数据缓存命中率直接跌倒谷底(为命中率 98%),各位看官执行分析其中原因。

g++ -O2 a.cpp -DRUN_LOOP2
valgrind --tool=cachegrind ./a.out

==1702874== Cachegrind, a cache and branch-prediction profiler
==1702874== Copyright (C) 2002-2017, and GNU GPL'd, by Nicholas Nethercote et al.
==1702874== Using Valgrind-3.18.1 and LibVEX; rerun with -h for copyright info
==1702874== Command: ./a.out
==1702874== 
--1702874-- warning: L3 cache found, using its data for the LL simulation.
^@s = 70
Time for loop2: 699.947ms
==1702874== 
==1702874== I   refs:      162,420,571
==1702874== I1  misses:          2,236
==1702874== LLi misses:          2,235
==1702874== I1  miss rate:        0.00%
==1702874== LLi miss rate:        0.00%
==1702874== 
==1702874== D   refs:       40,771,337  (40,567,464 rd   + 203,873 wr)
==1702874== D1  misses:     40,016,408  (40,013,784 rd   +   2,624 wr)
==1702874== LLd misses:     10,013,844  (10,012,087 rd   +   1,757 wr)
==1702874== D1  miss rate:        98.1% (      98.6%     +     1.3%  )
==1702874== LLd miss rate:        24.6% (      24.7%     +     0.9%  )
==1702874== 
==1702874== LL refs:        40,018,644  (40,016,020 rd   +   2,624 wr)
==1702874== LL misses:      10,016,079  (10,014,322 rd   +   1,757 wr)
==1702874== LL miss rate:          4.9% (       4.9%     +     0.9%  )

下面我们谈谈多维数据在一维内存中访问的另一个关键因素,步长。

步长解决的是如何将多维张量中的 (i, j) 或 (i, j, k, …) 索引映射到内存中的地址,比如在内存中从一个索引移动到另一个索引时,沿着某个特定维度需要跨越多少个元素(不一定是字节)。

行优先存储模式的步长,代码段5。

import torch

x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])

print("Tensor:")
print(x)
print("Strides:", x.stride())      # 输出: (3, 1)

# 行优先(Row-major)二维布局  
- 形状(Shape):`[2, 3]`(2 行,3 列)  
- 典型步长(Strides):`[3, 1]`

解释:
- 第一个维度(“行”维度)的步长是 `3`,意味着如果你从第 `i` 行移动到第 `i+1` 行,你需要在内存中跳过 `3` 个元素。
- 第二个维度(“列”维度)的步长是 `1`,意味着相邻列的元素在内存中也是相邻的。

列优先存储模式的步长,代码段6。

import numpy as np
import torch

#创建一个 NumPy 的列优先(column-major)数组
a = np.array([[1, 2, 3],
             [4, 5, 6]], order='F')  # 'F' 表示 Fortran order(列优先)


x = torch.from_numpy(a)
print(x)
print("Strides:", x.stride()) #输出:(1,2)


#列优先(Column-major)二维布局  
- 形状(Shape):`[2, 3]`(2 行,3 列)  
- 典型步长(Strides):`[1, 2]`

解释:
- 第一个维度(“行”维度)的步长是 `1`,意味着从第 `i` 行移动到第 `i+1` 行只需要在内存中前进一步。
- 第二个维度(“列”维度)的步长是 `2`,意味着从第 `j` 列跳转到第 `j+1` 列时,你需要跳过 `2` 个元素

再来看一个多维张量的例子,代码段7。

import torch
t = torch.arange(0, 24).reshape(1, 2, 3, 4)
print(t.stride())

torch.arange(0, 24) 创建了一个一维张量,包含从 0 到 23 的 24 个数字:
[ 0,  1,  2, ..., 23]
然后 .reshape(1, 2, 3, 4) 将其重塑为一个 **4D 张量**,其形状是:
(第0维: batch=1, 第1维: channel=2, 第2维: height=3, 第3维: width=4)
也就是说这个张量可以看作是一个 batch size 为 1、有 2 个通道、每个通道有 3 行 4 列的数据。

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]]]])

执行 print(t.stride()),输出为:
(24, 12, 4, 1)
这表示在内存中访问该张量时,每个维度上的“步长”分别是:
dim=0 (batch) 24
要跳到下一个 batch 需要移动 24 个元素(但因为 batch=1,实际上没用)
dim=1 (channel) 12
要跳到下一个通道需要移动 12 个元素
dim=2 (height/行) 4
要跳到下一行需要移动 4 个元素
dim=3 (width/列) 1
要跳到下一列只需要移动 1 个元素

基于上文的铺垫,我们继续探讨Pytorch中几个重要的张量变换操作,切片(slice)、转置(transpose)、广播(broadcast)。

切片,指从一个大的数据结构中提取出一部分连续的数据子集,代码段8。

import torch

# 创建一个张量
tensor = torch.tensor([[1, 2, 3],
                       [4, 5, 6],
                       [7, 8, 9]])

# 获取第二行
print(tensor[1, :])         # 输出: tensor([4, 5, 6])

# 获取第二列
print(tensor[:, 1])         # 输出: tensor([2, 5, 8])

# 获取子张量:前两行、前两列
print(tensor[:2, :2])       # 输出: tensor([[1, 2],
                           #                [4, 5]])

# 每隔一行 + 每隔一列
print(tensor[::2, ::2])     # 输出: tensor([[1, 3],
                           #                [7, 9]])

对一个张量(或矩阵)进行“切片”时,实际上并没有复制内存中的任何数据值。相反,这是通过“零拷贝”的方式,调整以下三个要素来创建一个新的视图(view)。

这三个关键要素是:

1. 起始偏移量(Base offset / Data pointer offset)
  • 表示新视图从原始内存块中的哪个位置开始。
  • 它是一个指向内存中某个元素的指针偏移量。
2. 形状(Shape)
  • 描述这个新视图的维度大小。
  • 例如:原张量是 3x4,切片后可能是 2x31x4
3. 步长(Strides)
  • 指明在每个维度上移动一个索引单位时,需要跨越多少个内存元素。
  • 通过调整步长,可以实现按行、列或其他任意模式访问数据。

具体的,假设有一个 3×4 的张量(即形状为 [3, 4]),在内存中是行优先存储:

[[A00, A01, A02, A03],
 [A10, A11, A12, A13],
 [A20, A21, A22, A23]]

对应的内存布局为一维数组:

[A00, A01, A02, A03, A10, A11, A12, A13, A20, A21, A22, A23]

如果我们做如下切片操作(取第1行到第3行,第1列到第3列):

sub_tensor = tensor[1:, 1:3]

结果是一个形状为 [2, 2] 的新视图:

[[A11, A12],
 [A21, A22]]

但此时并没有复制任何数据,而是通过以下方式创建了一个新的“视图”,数据还是指向内存中已有的数据。

属性说明
数据指针偏移指向 A11 的位置即内存地址偏移 5(从 A00 开始数)
形状(shape)[2, 2]表示两行两列
步长(strides)[4, 1]行步长为 4(跳过一行),列步长为 1(逐列)

其下标索引公式可表示为:

$$A[i, j] = A.data[offset + i*A.strides[0] + j * A.strides[1]]$$

读者可自行验证。

转置,可以理解为对多维张量的各个维度(轴)进行重新排列。

比如,对于一个二维矩阵(形状为 `[m, n]`),最简单的转置就是将行和列交换,得到一个新的形状为 `[n, m]` 的矩阵。对于更高维的张量,”转置” 通常意味着对各个轴进行更一般的排列组合。例如,将一个形状为 `[D0, D1, D2]` 的张量转置为 `[D2, D0, D1]`。

转置包括逻辑转置物理转置

  1. 基于步长的逻辑转置(零拷贝

我们可以通改变对内存中数据缓冲区的解读方式来实现逻辑上的转置,而不是真正地重新排列数据。比如,在一个简单的二维情况(行优先布局)中:

  • 原始的步长可能是 strides = [n, 1],对应一个形状为 [m, n] 的数组。
  • 当进行转置时,只需将这些步长交换为 [1, m],就可以“按转置顺序”来读取数据。
  • 这种方法不需要复制数据,但要注意的是,如果原始数据是按行优先方式存储的,这种转置后的布局在缓存访问上可能效率较低。

具体的,代码段9。

import torch

# 创建一个张量
t = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])

print("Original tensor:")
print(t)
print("Strides:", t.stride())  # 输出: (3, 1)

# 逻辑转置
t_trans = t.t()  # 或者使用 t.transpose(0, 1)

print("\nTransposed view:")
print(t_trans)
print("Strides after transpose:", t_trans.stride())  # 输出: (1, 3)
  1. 物理转置(拷贝)

实际创建一个新的数据缓冲区,并将元素按照转置后的位置写入新内存,使转置后的数据在新的布局中是连续存储的。虽然这种方式需要花费时间和内存来完成拷贝,但它能为后续许多针对转置数据的操作带来更好的内存局部性(memory locality),比如在 TensorRT、ONNX Runtime等加速库中可以更好地利用缓存从而提升性能。

具体的,代码段10。

import torch

# 创建一个张量
t = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])
# 创建一个连续的转置副本
t_trans_contig = t.t().contiguous()

print("\nPhysically transposed and contiguous tensor:")
print(t_trans_contig)
print("Strides after contiguous:", t_trans_contig.stride())  # 输出: (2, 1)
print("Is contiguous?", t_trans_contig.is_contiguous())  # 应为 True

值得注意的是,逻辑转置中,转置后的矩阵步长是(1,3),而物理转置中,转置后的步长是(2,1)。

这里解释一下,前提条件是行优先存储,在逻辑转置中,转置前

[[1, 2, 3],
 [4, 5, 6]]

底层数据存储为[1,2,3,4,5,6],

转置后

[[1, 4],
 [2, 5],
 [3, 6]]

底层内存中的数据并没有被重新排列!还是:[1, 2, 3, 4, 5, 6],只是访问方式变了 —— 现在我们按“列优先”的方式去读这个数组,所以正确的strides是(1,3)。

行 i列 j对应值
001
014
102
115
203
216

而在物理转置中,由于使用了.contiguous(),内存空间得以连续布局,数据在底层存储为[1, 4, 2, 5, 3, 6],这个新内存布局使得每一行是连续的,因此步长为 (2,1)。

我们可以把 .contiguous() 理解为:“我不管你现在怎么解读这块内存,我现在要把你的数据重新整理成行优先、连续存储的方式。”

广播,是在张量(或数组)运算中非常常见且强大的一种机制,它允许不同形状的张量进行逐元素(element-wise)操作,而不需要显式复制数据。

我们使用了一个示例来说明这一概念:对两个张量 ab 进行相加操作,其中 a.shape = [4, 3]b.shape = [1, 3],如图 1 所示

该图展示了 广播(Broadcast) 在数组运算中的底层实现原理,重点在于 如何通过调整步幅(strides)实现数据虚拟扩展,而非物理复制数据。

  • 数组 a:形状为 4×3,假设为常规行优先存储(如 strides=[3,1])。
  • 数组 b:原始形状为 1×3,数据内容 [1, 2, 3],初始步幅 strides=[3,1]
  • b1×3虚拟扩展4×3,使其能与 a 进行逐元素运算(如加法)。将维度大小为 1 的轴扩展至目标大小(此处将行维度从 1 扩展到 4)。
  • 广播通过 设置步幅为 0 实现虚拟扩展:
  • 调整后的 b 属性
    • shape = [4, 3],扩展后的逻辑形状。
    • data = [1, 2, 3],物理数据未改变。
    • strides = [0, 1]行维度步幅 0:表示在行方向移动时,内存地址不变(复用同一行数据),列维度步幅 1:与原始步幅一致,按元素大小(假设为 1 字节)步进。
  • 扩展后的 b 在逻辑上表现为每行都是 [1, 2, 3],但物理内存中数据不复制。根据公式 $$b[i, j] = b.data[offset + i*strides[0] + j*strides[1]$$,

对任意行 i(0 ≤ i <4),代入调整后的步幅 strides=[0, 1],所有行均指向原始数据的第 j 个元素,实现虚拟复制。

参考文献

[1]Efficient PyTorch: Tensor Memory Format Matters