NumPy 是 Python 科学计算的基础库,核心是多维数组 ndarray。PyTorch 的张量在接口和思维方式上几乎完全照搬 NumPy,所以先把数组的形状、广播、索引和矩阵乘法弄熟,后面读模型代码会轻松很多。
以下代码都用 import numpy as np。
数组与形状
数组的三个基本属性:shape(每一维的长度)、dtype(元素类型)、ndim(维数)。
a = np.array([1, 2, 3])
b = np.array([[1.0, 2.0], [3.0, 4.0]])
print(a.shape, a.dtype) # (3,) int64
print(b.shape, b.dtype) # (2, 2) float64
print(b.ndim) # 2常用的创建方式:
np.zeros((2, 3)) # 全 0
np.ones((2, 3), dtype=np.float32) # 全 1,指定类型
np.arange(6).reshape(2, 3) # [[0 1 2] [3 4 5]]
rng = np.random.default_rng(0) # 推荐的随机数生成器写法
rng.normal(size=(2, 3)) # 标准正态分布在深度学习里,看代码时最重要的习惯是时刻知道每个数组的形状。一批文本的嵌入通常是 (batch, seq_len, hidden),一批图片是 (batch, channel, height, width)。
索引与切片
x = np.arange(12).reshape(3, 4)
x[0] # 第 0 行:[0 1 2 3]
x[:, 1] # 第 1 列:[1 5 9]
x[1:, ::2] # 从第 1 行起,每隔一列取:[[4 6] [8 10]]
x[x > 5] # 布尔索引:[6 7 8 9 10 11]改变形状:
x.T.shape # 转置:(4, 3)
x.reshape(4, 3) # 元素总数不变即可
x.reshape(-1) # -1 表示「这一维自动推算」,这里展平成 (12,)广播
两个形状不同的数组做逐元素运算时,NumPy 会按规则自动「扩展」较小的那个,这就是广播(broadcasting)。规则是:从最后一维往前对齐,每一维要么相等,要么其中一个是 1(缺失的维度视为 1)。
m = np.arange(6).reshape(2, 3) # (2, 3)
m + 10 # 标量广播到每个元素
m + np.array([1, 2, 3]) # (3,) → 每一行都加 [1 2 3]
m + np.array([[100], [200]]) # (2, 1):第 0 行加 100,第 1 行加 200
# [[100 101 102]
# [203 204 205]]神经网络里给每个样本加上同一个偏置向量 x @ W + b,靠的就是广播。
矩阵乘法
@ 是矩阵乘法运算符,等价于 np.matmul。形状规则是 (m, k) @ (k, n) → (m, n),中间那一维必须相等。
A = rng.normal(size=(2, 3))
B = rng.normal(size=(3, 4))
(A @ B).shape # (2, 4)
batch = rng.normal(size=(8, 2, 3))
(batch @ B).shape # (8, 2, 4):前面的维度当作批次,逐个相乘注意 * 是逐元素相乘,不是矩阵乘法,这是初学者最常见的错误之一。矩阵乘法的数学含义见线性代数。
沿轴归约
sum、mean、max 等函数可以用 axis 指定沿哪一维计算。keepdims=True 保留被归约的那一维(长度变成 1),方便接着做广播。
x = np.arange(12).reshape(3, 4)
x.sum() # 66
x.sum(axis=0) # 沿行方向加,得到每列之和:[12 15 18 21]
x.sum(axis=1) # 每行之和:[ 6 22 38]
x.mean(axis=1, keepdims=True).shape # (3, 1)一个综合例子是数值稳定的 softmax:先减去每行最大值,避免 exp 溢出,再归一化。
def softmax(z):
z = z - z.max(axis=-1, keepdims=True)
e = np.exp(z)
return e / e.sum(axis=-1, keepdims=True)
softmax(np.array([[1.0, 2.0, 3.0], [1000.0, 1000.0, 1000.0]]))
# [[0.09 0.245 0.665]
# [0.333 0.333 0.333]] 第二行不减最大值的话会溢出成 nan向量化
Python 的 for 循环很慢,NumPy 的数组运算在底层用 C 实现并批量处理。能用数组运算表达的,就不要写循环,这叫向量化(vectorization)。
v = rng.normal(size=1_000_000)
s = 0.0
for e in v: # 纯 Python 循环
s += e * e
s = np.dot(v, v) # 向量化:结果相同,快几十到上百倍同样的思路贯穿整个深度学习:GPU 擅长的正是大规模的批量矩阵运算。
常见 dtype
| dtype | 说明 |
|---|---|
int64 | 常用于 token ID、类别标签;NumPy 默认整数是平台相关的 np.int_,在常见的 64 位 NumPy 2.x 环境中为 int64 |
float64 | 浮点数组的默认类型,NumPy 默认用它 |
float32 | 深度学习的传统默认精度,PyTorch 默认用它 |
float16 | 半精度,最大只能表示约 65504,容易溢出 |
bool | 布尔掩码 |
np.array([1, 2]).dtype # int64
np.array([1.0]).dtype # float64
np.array([1.5, 2.7]).astype(np.int32) # [1 2]:直接截断,不是四舍五入NumPy 默认 float64,PyTorch 默认 float32,两者互相转换时要留意精度。NumPy 本身不支持 bf16,低精度训练的细节见PyTorch张量与自动求导。
在 AI 中的作用
- 数据预处理、评估指标计算、结果分析大量使用 NumPy。
- PyTorch 的张量 API 与 NumPy 高度一致,形状、广播、索引的规则基本通用。
- 理解广播和矩阵乘法的形状规则,是读懂注意力机制等模型代码的前提。
相关
- 线性代数:向量、矩阵与矩阵乘法的数学含义。
- PyTorch张量与自动求导:在 GPU 上运行、可以自动求导的「NumPy」。
- 注意力机制:大量使用批量矩阵乘法、softmax 与广播。