PyTorch张量操作

PyTorch 的 Tensor 表面上像一个多维数组：可以切片、转置、相加、矩阵乘、放到 GPU 上计算。真正写模型或调试 shape 问题时，更关键的是另一层问题：这次操作有没有复制数据，新的 shape 是怎么解释同一段内存的，为什么有些 tensor 明明形状对却不能 view()，为什么 GPU tensor 的 .shape 可以直接打印而 .item() 会让 CPU 等待。

把 Tensor 理解成“数据缓冲区 + 元信息 + 一套索引解释规则”，很多 API 的行为会变得自然很多。

1
import torch
2

3
torch.manual_seed(0)

1. Tensor 内部由什么组成#

普通 dense tensor 不是一个嵌套的 Python list，也不是每一行各自分配一块内存。更接近实际情况的模型是三层：

Python 侧的 torch.Tensor 对象，是用户代码拿到的句柄。
底层 C++ 的 TensorImpl 记录这个 tensor 的元信息，例如 size、stride、 storage offset、dtype、device、layout，以及 autograd 和版本计数相关状态。
StorageImpl / DataPtr 指向实际数据缓冲区，并记录分配器、字节数和释放方式。

Tensor 内部结构

对普通 torch.strided dense tensor 来说，数值数据通常是一段扁平连续的 byte buffer。 dtype 决定每个元素占多少字节，shape 和 stride 决定多维下标如何映射到这段 buffer。这里有一个很重要的区分：storage 通常是一段连续分配的底层缓冲区，但 tensor 这个逻辑视图不一定是 contiguous 的。

GPU tensor 也遵循类似的对象模型，只是实际数值 buffer 位于设备内存中。Python 对象、 TensorImpl、shape、stride、dtype、device、layout、引用计数和大部分调度元信息仍由宿主进程里的 PyTorch 运行时管理。读取 x.shape、x.dtype、x.stride() 不需要把整块 GPU 数据复制回 CPU；但读取真实数值，例如 x.cpu()、x.item()、print(x)，就必须让 CPU 观察设备上的数据，通常会触发数据传输或同步。

1
def show(name, x):
2
    print(name)
3
    print("  shape:", tuple(x.shape))
4
    print("  stride:", x.stride())
5
    print("  storage_offset:", x.storage_offset())
6
    print("  contiguous:", x.is_contiguous())
7
    print("  dtype/device:", x.dtype, x.device)
8
    print("  element_size:", x.element_size())
9
    print("  numel:", x.numel())
10
    print("  storage_nbytes:", x.untyped_storage().nbytes())
11
    print("  data_ptr:", hex(x.data_ptr()))
12
    print("  storage_ptr:", hex(x.untyped_storage().data_ptr()))
13

14

15
x = torch.arange(6, dtype=torch.int64).reshape(2, 3)
16
y = x[:, 1:]
17
z = x.t()
18

19
show("x", x)
20
show("y = x[:, 1:]", y)
21
show("z = x.t()", z)

y 和 z 都没有重新分配一份完整数据。它们复用 x 的 storage，只是记录了不同的 shape、stride 和 storage_offset。data_ptr() 表示当前 tensor 第一个逻辑元素的地址，切片后可能会变；untyped_storage().data_ptr() 表示底层 storage 的起点，view 之间通常相同。

2. Storage、shape 和 stride#

以 x = torch.arange(6).reshape(2, 3) 为例，底层 storage 可以想象成：

1
storage: 0  1  2  3  4  5

x 的逻辑形状是 (2, 3)，stride 是 (3, 1)。访问 x[i, j] 时，PyTorch 用下面的公式找到 storage 中的元素位置：

1
storage_index = storage_offset + i * stride[0] + j * stride[1]

所以 x[1, 2] 对应 0 + 1 * 3 + 2 * 1 = 5，也就是 storage 中的第 5 号元素。

Storage 与 stride

转置不会移动数据。x.t() 的 shape 变成 (3, 2)，stride 变成 (1, 3)。逻辑上的 z[2, 1] 仍然落到 0 + 2 * 1 + 1 * 3 = 5，读到的仍是同一段 storage 里的元素。这就是 view 操作便宜的原因：很多时候只改元信息，不搬运数据。

contiguous 描述的是逻辑 tensor 按当前 shape 访问时，元素是不是按紧密顺序排列。原始 x 是 contiguous 的，x.t() 通常不是：

1
x = torch.arange(6).reshape(2, 3)
2
z = x.t()
3

4
print(x.shape, x.stride(), x.is_contiguous())
5
print(z.shape, z.stride(), z.is_contiguous())

storage 连续不等于 tensor contiguous。一个 non-contiguous view 仍然可以引用一段连续 storage，只是它的逻辑相邻元素在 storage 中不一定相邻。

3. 创建张量#

创建 tensor 时要先分清楚：已有数据要变成 tensor，还是只需要按某个形状分配一块新 tensor。

1
# 从 Python 数据创建，通常会复制出新的 tensor 数据。
2
a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
3

4
# 只按形状创建。
5
b = torch.empty(2, 3)          # 未初始化，内容不可假设
6
c = torch.zeros(2, 3)
7
d = torch.ones(2, 3)
8
e = torch.full((2, 3), 7)
9

10
# 按数列创建。
11
r = torch.arange(0, 10, 2)
12
l = torch.linspace(0, 1, steps=5)
13

14
# 沿用已有 tensor 的 dtype/device 等属性。
15
same_shape = torch.zeros_like(a)
16
different_shape = a.new_ones(3, 4)

torch.empty 只分配内存，不清零。它适合马上被完整写入的场景，不适合直接参与计算：

1
x = torch.empty(3)
2
print(x)  # 里面是未初始化值，不能当作随机数或 0 使用

从 NumPy 或已有对象转 tensor 时，是否共享底层内存也很重要：

1
import numpy as np
2

3
arr = np.array([1, 2, 3], dtype=np.float32)
4

5
x = torch.tensor(arr)       # 复制
6
y = torch.as_tensor(arr)    # 能共享时通常共享
7
z = torch.from_numpy(arr)   # 与 NumPy array 共享 CPU 内存
8

9
arr[0] = 100
10
print(x)  # tensor([1., 2., 3.])
11
print(y)  # tensor([100., 2., 3.])
12
print(z)  # tensor([100., 2., 3.])

旧式构造器 torch.Tensor(...) 容易和 dtype 默认行为混在一起，新代码优先使用 torch.tensor(...)、torch.empty(...)、torch.zeros(...) 这类明确的创建函数。

4. dtype 和 device#

每个 tensor 都有 dtype、device 和 layout。最常见的 layout 是 torch.strided。 dtype 决定元素解释方式和单个元素大小，device 决定数据 buffer 分配在哪类设备上。

1
x = torch.arange(6)
2
print(x.dtype, x.device, x.layout)
3

4
xf = x.to(dtype=torch.float32)
5
print(xf.dtype)
6

7
if torch.cuda.is_available():
8
    xc = xf.to("cuda")
9
    print(xc.device)
10
    print(xc.shape, xc.stride())  # 读元信息，不需要复制整块数据回 CPU

CPU tensor 和 GPU tensor 不会在普通算子里自动混算：

1
cpu_x = torch.ones(3)
2

3
if torch.cuda.is_available():
4
    gpu_y = torch.ones(3, device="cuda")
5

6
    # RuntimeError: expected all tensors to be on the same device
7
    # cpu_x + gpu_y
8

9
    z = cpu_x.to(gpu_y.device) + gpu_y
10
    print(z.device)

不同 dtype 混合时，PyTorch 会做 dtype promotion。结果 dtype 往往能容纳更多信息，但 in-place 写回时仍受左侧 tensor dtype 限制。

1
i = torch.ones(3, dtype=torch.int32)
2
f = torch.ones(3, dtype=torch.float32) * 0.5
3

4
print((i + f).dtype)  # torch.float32
5

6
try:
7
    i.add_(f)
8
except RuntimeError as e:
9
    print(type(e).__name__, str(e).splitlines()[0])

to() 是 dtype/device 转换的统一入口。x.float()、x.long()、x.cpu()、x.cuda() 只是更短的常用写法。

5. 形状变换#

形状变换可以分成三类：只改解释方式、可能复制、一定创建新数据。最容易混淆的是 view() 和 reshape()。

view() 要求当前 tensor 的 stride 能直接按新 shape 解释。contiguous tensor 上通常没问题：

1
x = torch.arange(12).reshape(3, 4)
2
y = x.view(2, 6)
3

4
print(x.shape, x.stride())
5
print(y.shape, y.stride())

转置后的 tensor 通常不是 contiguous，直接 view() 会失败：

1
x = torch.arange(12).reshape(3, 4)
2
t = x.t()
3

4
print(t.shape, t.stride(), t.is_contiguous())
5

6
try:
7
    t.view(12)
8
except RuntimeError as e:
9
    print(type(e).__name__, str(e).splitlines()[0])

reshape() 更宽松。它能返回 view 时就返回 view，不能时会复制出一块符合目标 shape 的新 tensor。代码不应该依赖 reshape() 到底有没有复制。

1
r = t.reshape(12)
2
print(r.shape, r.is_contiguous())

如果确实需要连续布局，可以显式使用 contiguous()：

1
c = t.contiguous()
2
v = c.view(12)
3

4
print(c.shape, c.stride(), c.is_contiguous())

其他常见形状操作：

1
x = torch.randn(2, 1, 3)
2

3
print(x.squeeze(1).shape)       # torch.Size([2, 3])
4
print(x.unsqueeze(0).shape)     # torch.Size([1, 2, 1, 3])
5

6
y = torch.randn(2, 3, 4)
7
print(y.transpose(1, 2).shape)  # torch.Size([2, 4, 3])
8
print(y.permute(2, 0, 1).shape) # torch.Size([4, 2, 3])
9
print(torch.movedim(y, 0, -1).shape)

transpose()、permute()、movedim() 通常只是改变维度解释，常常返回 view，并可能制造 non-contiguous tensor。后面如果要调用要求连续布局的操作，就需要检查 is_contiguous()。

6. View、copy 和 in-place#

view 共享底层 storage，copy 拥有独立 storage。修改 view 会影响 base tensor，修改 copy 不会。

View 与 copy 分类

1
base = torch.arange(6).reshape(2, 3)
2

3
v = base[:, 1:]
4
v.zero_()
5
print(base)
6

7
base = torch.arange(6).reshape(2, 3)
8
c = base[:, 1:].clone()
9
c.zero_()
10
print(base)

常见操作可以粗略分成三类：

类型	常见操作	判断方式
通常返回 view	basic slicing、`view`、`transpose`、`permute`、`squeeze`、`unsqueeze`、`expand`	不搬运数据，改元信息
通常返回 copy	`clone`、advanced indexing、`repeat`	新 storage
视情况而定	`reshape`、`flatten`、`contiguous`	能复用就复用，必要时复制

带下划线的方法通常表示 in-place mutation：

1
x = torch.tensor([1.0, 2.0, 3.0])
2
y = x
3

4
x.add_(10)
5
print(y)  # y 和 x 指向同一个 tensor 对象，能看到修改

索引赋值、copy_()、许多带 out= 的函数也会写入已有 tensor：

1
x = torch.zeros(5)
2
x[1:4] = torch.tensor([1.0, 2.0, 3.0])
3
print(x)
4

5
out = torch.empty(3)
6
torch.add(torch.ones(3), 2, out=out)
7
print(out)

detach() 和 clone() 解决的是不同问题：

1
x = torch.ones(3, requires_grad=True)
2

3
a = x.detach()          # 切断 autograd 历史，但通常仍共享 storage
4
b = x.clone()           # 复制数据，仍保留 autograd 关系
5
c = x.detach().clone()  # 切断 autograd 历史，并复制数据

训练代码里不要随意对需要梯度的中间结果做 in-place 修改。某些中间值会被 autograd 保存给 backward 使用，原地改掉以后，反向传播可能报版本计数错误，或者让代码语义变得难以推断。

7. 索引、切片、gather 和 scatter#

basic indexing 包括整数索引、切片、...、None。这类索引通常返回 view：

1
x = torch.arange(12).reshape(3, 4)
2

3
v = x[1:, 1:3]
4
v[0, 0] = -1
5

6
print(x)

advanced indexing 包括整数列表、LongTensor index、bool mask 等。读取时通常返回 copy：

1
x = torch.arange(12).reshape(3, 4)
2

3
y = x[[0, 2], [1, 3]]
4
y[0] = -99
5

6
print(y)
7
print(x)  # x 不变

索引赋值本身是 in-place，不管索引形式是 basic 还是 advanced：

1
x = torch.arange(12).reshape(3, 4)
2
mask = x % 2 == 0
3
x[mask] = 0
4
print(x)

index_select() 适合沿某个维度抽取一组位置：

1
x = torch.arange(12).reshape(3, 4)
2
idx = torch.tensor([0, 2])
3

4
print(torch.index_select(x, dim=0, index=idx))
5
print(torch.index_select(x, dim=1, index=idx))

gather() 按 index tensor 从源 tensor 中取值。index 的 shape 决定输出 shape，index 中的数值表示沿 dim 这一维要取源 tensor 的哪个位置。

1
scores = torch.tensor([
2
    [0.1, 0.8, 0.1],
3
    [0.2, 0.3, 0.5],
4
])
5

6
index = torch.tensor([
7
    [1],
8
    [2],
9
])
10

11
picked = torch.gather(scores, dim=1, index=index)
12
print(picked)

scatter_() 则是反方向：按照 index 把值写回目标 tensor。

1
out = torch.zeros(2, 3)
2
src = torch.tensor([
3
    [9.0],
4
    [8.0],
5
])
6

7
out.scatter_(dim=1, index=index, src=src)
8
print(out)

如果多个位置写到同一个目标位置，普通 scatter_() 的语义不适合表达累加，应该使用 scatter_add_() 或 scatter_reduce_() 这类明确的归约写入。

8. Broadcasting#

broadcasting 的规则从最后一维开始向左对齐。每一维要么相等，要么其中一个是 1，要么其中一个 tensor 没有这一维。结果 shape 在每一维上取较大的那个。

1
x = torch.empty(5, 1, 4, 1)
2
y = torch.empty(3, 1, 1)
3

4
print((x + y).shape)  # torch.Size([5, 3, 4, 1])

Broadcasting 对齐规则

expand() 是逻辑扩展，不复制数据；repeat() 会真实复制数据。

1
v = torch.tensor([1, 2, 3]).view(3, 1)
2

3
e = v.expand(3, 4)
4
r = v.repeat(1, 4)
5

6
print(e)
7
print("expand stride:", e.stride())
8
print("repeat stride:", r.stride())
9
print("expand storage bytes:", e.untyped_storage().nbytes())
10
print("repeat storage bytes:", r.untyped_storage().nbytes())

expand 后被扩展的维度通常会出现 stride 为 0 的情况。逻辑上有多个元素，物理上可能重复读同一个内存位置。对这种 view 做 in-place 写入很危险，PyTorch 会在许多场景下直接拒绝。

in-place broadcasting 还有一个限制：左侧 tensor 的 shape 不能因为广播而改变。

1
x = torch.empty(1, 3, 1)
2
y = torch.empty(3, 1, 7)
3

4
try:
5
    x.add_(y)
6
except RuntimeError as e:
7
    print(type(e).__name__, str(e).splitlines()[0])

报 shape mismatch 时，先把 shape 从右向左对齐，通常比盯着报错信息更快。

9. 逐元素、规约和矩阵运算#

逐元素操作会在相同位置上计算，必要时先按 broadcasting 规则对齐 shape：

1
x = torch.randn(2, 3)
2
y = torch.randn(3)
3

4
z = torch.clamp(torch.exp(x + y), max=10)
5
print(z.shape)

比较和逻辑操作返回 bool tensor：

1
x = torch.tensor([-1.0, 0.0, 2.0])
2
mask = x > 0
3
print(mask, mask.dtype)
4
print(torch.where(mask, x, torch.zeros_like(x)))

规约操作最重要的是 dim 和 keepdim：

1
x = torch.randn(2, 3, 4)  # batch, seq, hidden
2

3
print(x.sum().shape)
4
print(x.sum(dim=1).shape)
5
print(x.sum(dim=1, keepdim=True).shape)
6

7
values, indices = x.max(dim=-1)
8
print(values.shape, indices.shape)

keepdim=True 会保留被规约的维度，长度变成 1。这样后续和原 tensor 做 broadcasting 会更自然：

1
x = torch.randn(2, 3, 4)
2
mean = x.mean(dim=-1, keepdim=True)
3
centered = x - mean
4
print(centered.shape)

矩阵乘法相关 API 的 shape 约定不同：

1
a = torch.randn(3, 4)
2
b = torch.randn(4, 5)
3
print(torch.mm(a, b).shape)  # 只处理二维矩阵
4

5
ba = torch.randn(2, 3, 4)
6
bb = torch.randn(2, 4, 5)
7
print(torch.bmm(ba, bb).shape)  # 批量二维矩阵乘
8

9
w = torch.randn(4, 5)
10
print(torch.matmul(ba, w).shape)  # 支持 broadcasting 的通用 matmul
11
print((ba @ w).shape)

einsum 适合表达更复杂的维度关系，但不应该为了炫技替代清晰的 matmul、sum 或 transpose：

1
x = torch.randn(2, 3, 4)
2
w = torch.randn(4, 5)
3

4
print(torch.einsum("bsh,hk->bsk", x, w).shape)

10. 拼接、拆分和重排#

cat 沿已有维度拼接，stack 新增一个维度再拼接：

1
a = torch.ones(2, 3)
2
b = torch.zeros(2, 3)
3

4
print(torch.cat([a, b], dim=0).shape)    # torch.Size([4, 3])
5
print(torch.cat([a, b], dim=1).shape)    # torch.Size([2, 6])
6
print(torch.stack([a, b], dim=0).shape)  # torch.Size([2, 2, 3])
7
print(torch.stack([a, b], dim=1).shape)  # torch.Size([2, 2, 3])

拆分操作常用于把一个大 tensor 切成若干块：

1
x = torch.arange(10)
2

3
print(torch.split(x, 3))
4
print(torch.chunk(x, 3))
5
print(torch.tensor_split(x, 3))

unbind 会沿某个维度去掉这一维，并返回多个 tensor：

1
x = torch.arange(6).reshape(2, 3)
2
rows = torch.unbind(x, dim=0)
3
cols = torch.unbind(x, dim=1)
4

5
print(rows)
6
print(cols)

其他常用重排操作：

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

3
print(torch.repeat_interleave(x, repeats=2))
4
print(torch.tile(x, dims=(2,)))
5
print(torch.roll(x, shifts=1))
6
print(torch.flip(x, dims=(0,)))

这些操作看起来都在“变形”，但内存行为不同。性能敏感路径中，repeat、tile、 advanced indexing、contiguous()、某些 reshape() 都可能带来真实复制。

11. 调试检查清单#

Tensor 相关 bug 大多可以先打印以下信息：

1
def debug_tensor(name, x):
2
    print(
3
        name,
4
        "shape=", tuple(x.shape),
5
        "stride=", x.stride(),
6
        "dtype=", x.dtype,
7
        "device=", x.device,
8
        "contiguous=", x.is_contiguous(),
9
    )

常用判断：

shape mismatch：把两个 shape 从右向左对齐，按 broadcasting 规则逐维检查。
view() 报错：检查 is_contiguous() 和 stride()，必要时用 reshape() 或 contiguous().view()。
结果被意外修改：检查是不是拿到了 view，或者是不是做了 in-place 操作。
显存或内存突然上涨：检查 repeat、advanced indexing、contiguous()、CPU/GPU 往返。
CUDA 程序突然变慢：检查循环里是否频繁 .item()、.cpu()、print(cuda_tensor)。

一个实用习惯是：只看数值不够，shape、stride、dtype、device 四个属性要一起看。很多 PyTorch 张量操作的差别不在数学结果，而在它们是否共享 storage、是否改变 contiguous 状态、是否把数据搬到了另一个 device。