NumPy 数组、向量化与广播

上一课你把环境搭好了：conda 环境、解释器、能跑代码的编辑器。现在地基要从"能运行"升级到"能高效地处理一整块数据"。本模块的终点是从零写出一个训练循环，而训练循环的每一步——前向、损失、反向——本质上都是在操作一个个多维数组。这一课我们建立"张量（tensor）心智模型"：数据是带形状（shape）和类型（dtype）的内存块，运算尽量整块地做、不写 Python 循环。这套心智模型会 1:1 迁移到下一课的 PyTorch 张量上——届时你只是把 np 换成 torch、再加上"设备"和"梯度"两个新维度而已。

ndarray：shape 与 dtype 是张量的"身份证"

一个 NumPy 数组（ndarray）= 一块连续内存缓冲区 + 一个描述如何解读它的"头部"（header）。头部里最重要的两项是 shape（每一维多长，写成 (N, D) 这样的元组）和 dtype（每个元素的类型，如 float64、int64）。同一块缓冲区，配不同的头部，就能被解读成不同形状——这正是 reshape 几乎免费的原因，也是下面"视图"的关键。

• shape 决定广播和矩阵乘法能不能对上。ML 里 99% 的 bug 是 shape 不对，养成随手 print(x.shape) 的习惯。
• dtype 决定精度与显存。深度学习默认 float32（而非 NumPy 默认的 float64），因为省一半显存、且对训练足够。整数索引必须是整型 dtype。

view vs copy：切片是视图，就地改会"串"到原数组

机制直觉：切片返回视图——一个新头部，指向同一块内存缓冲区。所以对切片就地赋值，改的是底层那块内存，原数组跟着变。而 .copy() 会另开一块缓冲区，互不影响。这不是 NumPy 的怪癖，而是性能设计：切一个亿级数组若每次都拷贝就太慢了。代价是你必须时刻清楚手里是视图还是拷贝。

先给一个判断技巧的直觉：每个数组都有个 .base 属性，它就是"我这块内存是从谁那借来的"的指针——借来的（视图）就指向原数组，自己新开内存的就是 None。

import numpy as np
a = np.arange(10)
b = a[1:]          # 切片 -> 视图，和 a 共享内存
b[0] = 99          # 就地改 b 的第 0 个元素
print("a =", a)              # a 也变了！a[1] 现在是 99
print("b is view:", b.base is a)   # True：b 的底层缓冲就是 a

# 注意：a 此刻已是 [0 99 2 3 ...]。下面演示 copy，用一个全新的数组避免混淆
a2 = np.arange(10)
c = a2[1:].copy()  # 显式拷贝 -> 新缓冲区
c[0] = -1
print("c base is None:", c.base is None)  # True：拷贝自己新开内存
print("a2 after copy edit:", a2)          # a2 不受影响：[0 1 2 3 ...]

运行结果：第一段里 a = [ 0 99 2 3 ... ]（视图被改，原数组跟着变）；而对 copy 出来的 c 就地改写，全新的 a2 纹丝不动。判断技巧：arr.base is not None 通常意味着它是某个数组的视图（注意 reshape 等也会令 base 非 None，所以是"通常"而非绝对）。基本切片返回视图；花式索引和布尔掩码返回拷贝（见后文）。这条规则到 PyTorch 完全一样——下一课你会看到 torch 也区分 view 与 contiguous。

import torch
a = torch.arange(10)
b = a[1:]          # 同样是视图，共享 storage
b[0] = 99
print(a)           # tensor([ 0, 99, 2, ...])
c = a[1:].clone()  # PyTorch 用 .clone() 而非 .copy()

向量化：把 for 循环交给底层 C/SIMD

机制直觉：Python 的 for 循环每次迭代都要做动态类型检查、对象拆箱，开销巨大。向量化是把整块数组交给 NumPy 内部的 C 实现一次性算完，那里数据是紧密排列的同类型机器数，可以用 SIMD（单指令多数据，Single Instruction Multiple Data，一条指令同时算一排数）并行处理。结果常常快一个数量级以上。规则很简单：看到对数组元素逐个循环做同一件事，就想办法用数组运算整体表达。

下面用欧氏距离做对比：对 (N, D) 的每一行算它到查询向量 q 的距离。

import numpy as np, time
rng = np.random.default_rng(0)
N, D = 2000, 64
X = rng.standard_normal((N, D)); q = rng.standard_normal(D)

t0 = time.perf_counter()
out_loop = [sum((X[i,j]-q[j])**2 for j in range(D))**0.5 for i in range(N)]
t1 = time.perf_counter()

out_vec = np.sqrt(((X - q)**2).sum(axis=1))   # 一行搞定，全程无 Python 循环
t2 = time.perf_counter()

loop_t, vec_t = t1-t0, t2-t1
print(f"loop  : {loop_t*1000:.1f} ms")
print(f"vec   : {vec_t*1000:.3f} ms")
print(f"speedup: {loop_t/vec_t:.0f}x")   # 不同机器/环境数字会变，关注的是数量级而非确切倍数
print("max diff:", np.max(np.abs(np.array(out_loop)-out_vec)))  # ~1e-15，结果一致

注意 X - q 这一步：X 是 (2000, 64)，q 是 (64,)，却能直接相减——这就是广播在起作用。max diff 在 1e-15 量级，说明向量化版本和循环版本数值上等价，只是快了一个数量级。你跑出来的具体倍数会随机器和环境（桌面 CPython、Pyodide 等）波动，跑出 7× 还是 300× 都正常，别因此怀疑自己写错——关注量级即可。规格里要求用 %timeit 的场景，我们一律用 time.perf_counter 手写计时（Pyodide 里没有 IPython 魔法命令）。

广播：尾部对齐规则

广播让不同 shape 的数组也能逐元素运算，省去显式复制数据。规则只有一条，从最后一维（最右）开始向左逐维比较：

• 两维相等 → 这一维 OK；
• 其中一维是 1 → 把那个 1 "拉伸"到另一个的长度（不真复制，只是重复读取）；
• 否则 → 报错，不能广播。
• 维数不同时，短的那个在左侧补 1，再按上面比较。

经典例子：(N,1) 与 (1,M) 对齐 → (N,M)，这是构造距离矩阵、外积的常用手法。

务必看清的隐形 bug：(N,1) - (N,) → (N,N)

这是新手最常踩、且不会报错的坑。你想做"逐元素相减"得到 (N,)，结果悄悄变成了 (N,N) 的矩阵，后续计算全错却不报异常。原因：把 (N,) 左侧补 1 变成 (1,N)，再和 (N,1) 对齐 → (N,N)。

import numpy as np
x = np.array([1.0, 2.0, 3.0]).reshape(3, 1)   # (3,1)，列向量
y = np.array([1.0, 2.0, 3.0])                  # (3,)，一维
print("x.shape", x.shape, "y.shape", y.shape)
print("(x - y).shape:", (x - y).shape)         # 期望 (3,) 却得到 (3,3)！
print(x - y)
# 修复：让两边维度一致，都压成 (3,)
print("correct:", x.ravel() - y)               # [0. 0. 0.]

防身习惯：相减/相加前先 print 两个操作数的 shape；想逐元素就保证两边 shape 完全相同（或刻意只让该广播的那一维为 1）。

布尔掩码、花式索引、axis 与 keepdims

布尔掩码：a[a>0] 用一个同形状的布尔数组挑出满足条件的元素，返回一维拷贝。花式索引：a[[0,2,4]] 用整数数组按下标取元素，也返回拷贝。两者都不是视图——这点和切片相反，要记牢。

axis 指定沿哪一维做归约（reduction）：A.sum(axis=1) 把每一行的列加起来，(3,4)→(3,)，那一维"塌缩"消失了。keepdims=True 让被归约的维度保留为长度 1：(3,4)→(3,1)。为什么要保留？因为接下来往往要把归约结果广播回原形状（如按行归一化），(3,1) 能和 (3,4) 干净对齐，而 (3,) 会触发上面的尾部对齐错误。

import numpy as np
A = np.arange(12).reshape(3, 4).astype(float)
s1 = A.sum(axis=1)                 # (3,)
s2 = A.sum(axis=1, keepdims=True)  # (3,1)
print("no keepdims:", s1.shape, "| keepdims:", s2.shape)
print("row-normalized OK:\n", A / s2)   # (3,4)/(3,1) -> 每行和为 1
try:
    print(A / s1)                  # (3,4)/(3,) -> 尾部对齐失败
except ValueError as e:
    print("broadcast error:", e)

v = np.array([-3, 5, -1, 8, 0])
print("mask v>0:", v[v > 0])       # [5 8]，布尔掩码
print("fancy idx:", v[[0, 2, 4]])  # [-3 -1 0]，花式索引

记忆口诀：沿 axis 归约后想广播回去，就加 keepdims。这条在 PyTorch 里写作 keepdim=True（少个 s），机制完全一致。

数值稳定：softmax 为什么要减最大值

softmax 把一组实数 \(z\) 变成概率分布：\(\mathrm{softmax}(z)_i = e^{z_i} / \sum_j e^{z_j}\)。直接套公式有个致命问题：float 的指数会溢出。exp(1000) 就是 inf，inf/inf = nan，整个输出全是 nan。

关键观察：给所有 \(z_i\) 减去同一个常数 \(c\)，softmax 的值完全不变（分子分母同乘 \(e^{-c}\) 约掉）。取 \(c=\max_i z_i\)，则最大的指数变成 exp(0)=1，绝不会上溢。这就是"减最大值"的全部道理。

import numpy as np
z = np.array([1000.0, 1001.0, 1002.0])
# 朴素版：直接 exp，溢出
naive = np.exp(z) / np.exp(z).sum()
print("naive :", naive)              # [nan nan nan]
# 稳定版：减最大值，结果不变但不溢出
zs = z - z.max()
stable = np.exp(zs) / np.exp(zs).sum()
print("stable:", stable)             # [0.090 0.245 0.665]
# logsumexp 思想：稳定地算 log(sum(exp(z)))
lse = z.max() + np.log(np.exp(z - z.max()).sum())
print("logsumexp:", lse)             # 1002.41，不溢出

logsumexp 是同一招的推广：要算 \(\log\sum_i e^{z_i}\)（交叉熵损失里到处都是它），就提出最大值 \(\log\sum_i e^{z_i} = c + \log\sum_i e^{z_i-c}\)，其中 \(c=\max z\)。先放下细节，记住"凡是 exp/log 大数，先把最大值提出来"这个动作即可。PyTorch 直接提供 torch.logsumexp 和数值稳定的 F.log_softmax，你不必自己写——但要知道它们在背后做了这件事。

import torch
import torch.nn.functional as F
z = torch.tensor([1000., 1001., 1002.])
print(F.softmax(z, dim=0))     # 稳定，内部已减最大值
print(torch.logsumexp(z, 0))   # tensor(1002.41)

调一调，观察现象

微任务 1：把视图改成拷贝，观察"串改"消失。 这是两次独立运行的对比：先按原样跑一次（看到 a = [0 99 2 3 4]，视图被串改），再把第 3 行从 b = a[1:] 改成 b = a[1:].copy() 跑第二次（看到 a = [0 1 2 3 4]，原数组没动）。预期现象：改成拷贝后 a 不再被改动。为什么：.copy() 另开缓冲区，b[0]=99 写的是新内存，与 a 无关。

import numpy as np
a = np.arange(5)
b = a[1:]          # 第二次运行时改成 a[1:].copy() 再对比
b[0] = 99
print("a =", a)    # 第一次(视图)看到 [0 99 2 3 4]；改成 .copy() 第二次看到 [0 1 2 3 4]

微任务 2：给 reshape 加/去一个 1，看广播结果从 (N,) 变 (N,N)。 把 y 从 (3,) 改成 y.reshape(1,3) 或把 x 改成 (3,)，观察结果 shape。预期现象：只要有一边是二维且另一边长度对不上，就得到 (3,3)；两边都压成 (3,) 才得逐元素的 (3,)。为什么：尾部对齐时 1 被拉伸。

import numpy as np
x = np.array([10., 20., 30.]).reshape(3, 1)
y = np.array([1., 2., 3.])          # 试试 .reshape(1,3) 或保持 (3,)
res = x - y
print("shape:", res.shape)          # (3,3)
print(res)

微任务 3：把 softmax 的输入整体加 500，看朴素版崩、稳定版稳。 改 z = base + np.array([0.,1.,2.]) 里的 base，从 0 调到 800。预期现象：base 大到一定程度朴素版变 nan，稳定版始终输出同一组概率。为什么：softmax 平移不变，减最大值消除了溢出。

import numpy as np
base = 800.0                         # 调成 0 / 50 / 800 各跑一次
z = base + np.array([0., 1., 2.])
naive = np.exp(z) / np.exp(z).sum()
zs = z - z.max()
stable = np.exp(zs) / np.exp(zs).sum()
print("naive :", naive)
print("stable:", stable)             # 始终 [0.090 0.245 0.665]

动手练习

1. 欧氏距离矩阵（向量化）。 给定 A 形状 (N, D)、B 形状 (M, D)，不写任何 Python 循环，算出距离矩阵 Dist 形状 (N, M)，其中 Dist[i,j] 是 A[i] 与 B[j] 的欧氏距离。提示：用 A[:, None, :] - B[None, :, :] 触发 (N,1,D) 与 (1,M,D) 广播得到 (N,M,D)，再沿最后一维 sum 后开方。和双重循环版对比 perf_counter 与最大误差。
2. 稳定 softmax（按行）。 写 softmax(Z)，输入 (N, K)，对每一行做 softmax，输出 (N, K) 且每行和为 1。关键点：Z.max(axis=1, keepdims=True) 和归一化都要用 keepdims=True，想清楚去掉会怎样。用 print(out.sum(axis=1)) 验证全是 1。
3. （可选）一维卷积（valid）。 给信号 x（长 L）和核 w（长 k），向量化算 valid 卷积，输出长 L-k+1。先写循环版当正确性基准——这步最重要，向量化只是加速。挑战的向量化思路：用双重广播构造一个 (L-k+1, k) 的滑动窗口下标矩阵，即 idx = np.arange(L-k+1)[:, None] + np.arange(k)[None, :]（行偏移 + 列内偏移），再用 x[idx] 取出所有窗口，最后与 w 相乘并沿最后一维求和。如果广播还不熟，先把循环版跑对就够了，向量化留作挑战。

掌握自检

• 我能说出：切片返回______（视图/拷贝），布尔掩码和花式索引返回______，并能用 .base 验证。
• 给我两个 shape，比如 (8,1,6) 和 (7,6)，我能口算出能否广播、结果 shape 是什么。
• 我能立刻指出 (N,1) - (N,) 的结果是 (N,N) 而非 (N,)，并说明怎么修。
• 我能解释为什么 A / A.sum(axis=1) 报错而 A / A.sum(axis=1, keepdims=True) 正确。
• 我能在不查资料的情况下写出数值稳定的 softmax，并说出"减最大值"为什么不改变结果。
• 我能把一个逐元素的 Python 循环改写成向量化版本，并量出加速比。