张量运算:深度学习的通用语言

张量运算:深度学习的通用语言

本节摘要:张量(Tensor)是数据与深度学习之间的共同语言——每一张图像、每一个句子、每一个梯度都流经它。矩阵只能描述两组事物之间的成对关系,而真实数据无法塞进二维:32 张 224×224 的 RGB 图是 4 维张量 (32, 3, 224, 224),12 个注意力头的自注意力是 (batch, heads, seq_len, head_dim)。本节从几何与存储双重视角,讲透张量的三大基石:形状(Shape)、步长(Strides)、广播(Broadcasting);给出从零实现的 Tensor 类(扁平存储 + 步长索引),支持 reshape、transpose、squeeze/unsqueeze、逐元素运算与归约;演示 einsum(Einstein 求和)这一万能张量运算符,用一行表达式写出点积、外积、矩阵乘、批量矩阵乘、注意力分数;最后把多头注意力从头到尾用 einsum 走一遍,证明掌握张量运算后,形状错误便从「玄学」变成可机械调试的小事。

对应原课程:Phase 01 · Lesson 12 · tensor-operations(原英文 phases/01-math-foundations/12-tensor-operations/docs/en.md)。前置:第 1、2、3 节(线性代数)。

学习目标

阅读完本节,你应当能够:

1. 从零实现一个 Tensor 类,含形状、步长、reshape、transpose 与逐元素运算。
2. 应用广播规则,在不拷贝数据的前提下对不同形状的张量做运算。
3. 写出 einsum 表达式完成点积、矩阵乘、外积与批量运算。
4. 逐层追踪多头注意力中每一步的张量形状。

一、问题与直觉

你搭了一个 Transformer,前向传播看起来很干净,一跑就报错:RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x768 and 512x768)。你盯着形状看,试着转置,现在它说 Expected 4D input (got 3D input);你加了个 unsqueeze,别的地方又崩了。

形状错误是深度学习代码里最常见的 bug。它们在概念上不难——每个算子都有一份形状契约——但它们会飞速繁殖。一个 Transformer 把几十个 reshape、transpose、broadcast 串在一起,一个轴搞错,错误就级联放大。更糟的是,有些形状错误根本不抛异常:沿错误的轴广播、在错误的轴上求和,会悄悄产出垃圾。

矩阵处理的是两组事物之间的关系。真实数据塞不进二维。你需要一个能推广到任意维数、且运算能跨维干净组合的数据结构——这就是张量。掌握它的运算,形状错误就变得轻而易举。

1.1 张量是什么

张量是一个具有统一数据类型的多维数组。维数称为秩(Rank/阶),每一维是一个轴(Axis),形状(Shape) 是列出各轴大小的元组。

元素总数 = 各轴大小之积。形状 (2, 3, 4) 含 2*3*4 = 24 个元素。

1.2 深度学习中的张量形状

不同数据按惯例映射到特定形状:

PyTorch 用 NCHW(通道在前),TensorFlow 默认 NHWC(通道在后)。布局不匹配会导致隐蔽的变慢或报错。

1.3 内存布局与步长

二维数组在内存里是一段一维字节序列。步长(Strides) 告诉你沿每个轴前进一格要跳过多少元素。

行主序(C order):  a b c d e f   步长 (3, 1)   ← 每行连续
列主序(F order):  a d b e c f   步长 (1, 2)   ← 每列连续
​

transpose 不搬数据,只交换步长,让张量变得非连续(non-contiguous)——某行的元素在内存里不再相邻。

1.4 广播规则

广播让你对不同形状的张量做运算而无需拷贝数据。规则:从右对齐,两维相等或其中一维为 1 即可兼容;维数少的在左侧补 1。

张量 A:     (8, 1, 6, 1)
张量 B:        (7, 1, 5)
补齐 B:     (1, 7, 1, 5)
结果:       (8, 7, 6, 5)
​

1.5 Einsum:万能张量运算

Einstein 求和给每个轴贴一个字母标签。只出现在输入、不出现在输出的轴被求和消掉;同时出现在输入和输出的轴被保留。

常用模式:i,i->(点积)、i,j->ij(外积)、ii->(迹)、ij->ji(转置)、bij,bjk->bik(批量矩阵乘)、bhtd,bhsd->bhts(注意力分数)。

一次收缩的计算量是所有索引大小(保留与求和)之积。对 bij,bjk->bik,B=32、I=128、J=64、K=128,需 32*128*64*128 = 33,554,432 次乘加。

二、从零实现

完整源码见 phases/01-math-foundations/12-tensor-operations/code/tensors.py。

2.1 张量存储与步长

张量存一段扁平数字加形状元数据;步长告诉索引逻辑如何把多维下标映射到扁平位置。

class Tensor:
    def __init__(self, data, shape=None):
        if isinstance(data, (list, tuple)):
            self._data, self._shape = self._flatten_nested(data)
        elif isinstance(data, np.ndarray):
            self._data = data.flatten().tolist()
            self._shape = tuple(data.shape)
        else:
            self._data = [data]; self._shape = ()
        if shape is not None:
            total = reduce(lambda a, b: a * b, shape, 1)
            if total != len(self._data):
                raise ValueError(f"无法把 {len(self._data)} 个元素 reshape 成 {shape}")
            self._shape = tuple(shape)
        self._strides = self._compute_strides(self._shape)

    @staticmethod
    def _compute_strides(shape):
        if len(shape) == 0: return ()
        strides = [1] * len(shape)
        for i in range(len(shape) - 2, -1, -1):
            strides[i] = strides[i + 1] * shape[i + 1]
        return tuple(strides)
​

对形状 (3, 4),步长为 (4, 1)——前进一行跳 4 个元素,前进一列跳 1 个。

2.2 reshape / squeeze / unsqueeze

reshape 改形状不改元素顺序,元素总数必须不变;-1 表示该维由其他维推断。squeeze 去掉大小为 1 的轴;unsqueeze 插入一维——这是广播的关键:偏置向量 (D,) 要加到 batch (B, T, D) 上,得先 unsqueeze 成 (1, 1, D)。

2.3 transpose / permute

transpose 交换两轴,permute 重排所有轴——这就是 NCHW 与 NHWC 互转的方法。转置或 permute 后张量在内存里非连续;PyTorch 的 view 对非连续张量会失败,要改用 reshape 或先 .contiguous()。

2.4 逐元素运算与归约

逐元素运算(加、乘、减)独立作用于每个元素并保持形状;归约(sum、mean、max)塌缩一或多个轴。CNN 的全局平均池化:(B, C, H, W).mean(axis=[2,3]) 得 (B, C);NLP 的序列均值池化:(B, T, D).mean(axis=1) 得 (B, D)。

2.5 用 NumPy 演示广播

activations = np.random.randn(4, 3)
bias = np.array([0.1, 0.2, 0.3])
result = activations + bias                     # (4,3)+(3,) -> (4,3)

images = np.random.randn(2, 3, 4, 4)
scale = np.array([0.5, 1.0, 1.5]).reshape(1, 3, 1, 1)
result = images * scale                          # 按通道缩放

a = np.array([1, 2, 3]).reshape(-1, 1)
b = np.array([10, 20, 30, 40]).reshape(1, -1)
outer = a * b                                    # (3,4) 外积
​

用广播算两两距离:把 (M, 2) reshpe 成 (M, 1, 2)、(N, 2) 成 (1, N, 2),相减、平方、沿末轴求和、开根,得 (M, N)。

2.6 einsum 模式

dot    = np.einsum("i,i->", a, b)               # 点积
matmul = np.einsum("ik,kj->ij", A, B)           # 矩阵乘
batch_mm = np.einsum("bij,bjk->bik", batch_A, batch_B)  # 批量矩阵乘
​

2.7 用 einsum 走通多头注意力

B, H, T, D = 2, 4, 8, 16
E = H * D
X = np.random.randn(B, T, E)
W_q = np.random.randn(E, E) * 0.02

Q = np.einsum("bte,ek->btk", X, W_q)            # 投影
Q = Q.reshape(B, T, H, D).transpose(0, 2, 1, 3) # 拆头 -> (B,H,T,D)

scores = np.einsum("bhtd,bhsd->bhts", Q, K) / np.sqrt(D)
weights = softmax(scores, axis=-1)
attn_output = np.einsum("bhts,bhsd->bhtd", weights, V)

concat = attn_output.transpose(0, 2, 1, 3).reshape(B, T, E)
output = np.einsum("bte,ek->btk", concat, W_o)   # 输出投影
​

每一步都是张量运算:投影(einsum 矩阵乘)、拆头(reshape+transpose)、注意力分数(einsum 批量矩阵乘)、加权和(einsum 批量矩阵乘)、并头(transpose+reshape)、输出投影(einsum 矩阵乘)。

三、框架对比

从零实现 vs NumPy

从零实现 vs PyTorch

import torch
t = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
t.shape             # torch.Size([2, 3])
t.stride()          # (3, 1)
t.is_contiguous()   # True
t.reshape(3, 2); t.unsqueeze(0); t.transpose(0, 1).contiguous()
torch.einsum("ik,kj->ij", A, B)
​

PyTorch 额外提供自动微分、GPU 支持与优化 BLAS 内核,但形状语义完全一致。理解了从零版本,PyTorch 的形状错误就变得可读。

每个神经网络层都是张量运算

四、可复用产物

• outputs/prompt-tensor-shapes.md:一份系统化的张量形状调试提示,含每个常见算子(matmul、broadcast、cat、Linear、Conv2d、BatchNorm、softmax)的决策表与修复查找表。
• outputs/prompt-tensor-debugger.md:一个分步调试提示,把报错信息和你手上的张量形状贴进任何 AI 助手,即可拿回精确修复。

源码见 phases/01-math-foundations/12-tensor-operations/code/tensors.py。

五、练习

1. (Easy) reshape 往返:取形状 (2, 3, 4) 的张量,依次 reshape 成 (6, 4)、(24,),再回到 (2, 3, 4),每步打印扁平数据验证元素顺序不变。
2. (Medium) 实现广播:给 Tensor 类加 broadcast_to(shape) 方法,把大小为 1 的维扩展到目标形状;再改 _elementwise_op 使其运算前自动广播。用 (3,1) 与 (1,4) 测试得 (3,4)。
3. (Hard) 从零写 einsum:实现 einsum(subscripts, *tensors),至少支持点积 i,i->、矩阵乘 ij,jk->ik、外积 i,j->ij、转置 ij->ji。解析下标串、识别收缩索引、遍历所有索引组合,并与 np.einsum 对比。
4. (Hard) 注意力形状追踪器:写一个函数,输入 batch_size、seq_len、embed_dim、num_heads,打印多头注意力每一步的精确形状(输入、Q/K/V 投影、拆头、注意力分数、softmax 权重、加权和、并头、输出投影),并与 demo_attention_einsum() 的输出对照验证。

本节要点回顾

1. 张量是统一类型的多维数组——秩是轴数,形状是各轴大小元组,元素总数是各轴大小之积。
2. 步长决定内存布局——行主序步长 (3,1) 表示每行连续;transpose 只换步长、不搬数据,张量变非连续。
3. 广播从右对齐:两维相等或一维为 1 即兼容,维数少的左侧补 1;bias (D,) + (B,T,D) 自动广播。
4. einsum 是万能运算符:输入有、输出没有的轴被求和;一行写出点积、外积、矩阵乘、批量矩阵乘。
5. 形状错误最常见但可机械调试——每个算子都有形状契约,逐层追踪即可定位。
6. NCHW vs NHWC 是布局约定,PyTorch 用前者、TensorFlow 默认后者,不匹配会静默变慢。
7. 每个神经网络层都是张量运算——线性层、注意力、归一化、softmax 都能写成 einsum 或逐元素+广播。
8. PyTorch 形状语义与从零实现一致——掌握后者,前者的报错就变得可读。

下一节,我们直面浮点数的漏的抽象——数值稳定性:为什么训练会突然 NaN,log-sum-exp 技巧如何拯救 softmax 与交叉熵,以及 bfloat16 为何在训练中胜过 float16。