""" ============================================================================= CS230 · 第 9 讲工具模块：softmax 和 Batch Norm 的 numpy 从零实现（Ng 的列约定） ============================================================================= 对应讲义：lectures/lecture09_tuning_batchnorm_frameworks.md（§2 Batch Norm、§3 softmax） ----------------------------------------------------------------------------- 【这个模块是什么】 ----------------------------------------------------------------------------- 本讲起代码改用 PyTorch。但在用框架之前，我们先用 numpy 把本讲的两个新"积木"亲手写一遍， 再和 PyTorch 逐数字对照（代码 02、03）。写过一遍，你才知道 nn.BatchNorm1d、nn.CrossEntropyLoss 这两行代码背后做了什么。 softmax 一族（§3）： softmax_naive 教科书写法 e^z / Σe^z（z 很大时会溢出 —— 用来演示问题） softmax 数值稳定写法：先减去每一列的最大值 logsumexp、log_softmax log Σ e^z 的稳定写法；log softmax = z − logsumexp(z) cross_entropy_naive 先 softmax 再取 log（正确类的概率下溢成 0 时得到 inf） cross_entropy_from_logits 直接从 logits 算交叉熵（PyTorch nn.CrossEntropyLoss 的做法） softmax_jacobian 一个样本的 ∂ŷ/∂z（C×C 矩阵） softmax_model_forward / softmax_model_backward 在第 4 讲 deep_nn 的积木上换一个 softmax 输出层：dZ[L] = Ŷ − Y Batch Norm 一族（§2）： batchnorm_forward Z → Ẑ（减均值、除标准差）→ Z̃ = γ Ẑ + β batchnorm_backward 紧凑的反向公式（讲义 §2.4 推导） batchnorm_backward_staged 按计算图一步一步用链式法则（和紧凑公式应当逐数字一致） ema_update 测试时用的"指数加权平均"（running mean / var） ----------------------------------------------------------------------------- 【形状约定】和第 1–8 讲一样：**每一列一个样本** ----------------------------------------------------------------------------- logits Z、概率 Ŷ、one-hot 标签 Y ：(C, m)，C 是类别数 BN 的输入 Z、输出 Z̃ ：(n[l], m) γ、β、μ、σ² ：(n[l], 1)，和第 4 讲的 b[l] 同形状 PyTorch 的习惯正好相反（每一行一个样本），所以和 PyTorch 对照时要转置（.T）—— 讲义 §4.3。 【依赖】numpy；softmax_model_* 只读地复用第 4 讲 code/lecture04/deep_nn.py 的积木 （linear_activation_forward、linear_backward、linear_activation_backward）。 【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【numpy 小知识】【Python 小知识】 """ import os import sys import numpy as np _HERE = os.path.dirname(os.path.abspath(__file__)) for _d in ("lecture01", "lecture04"): _p = os.path.join(_HERE, "..", _d) if _p not in sys.path: sys.path.insert(0, _p) import deep_nn # noqa: E402 （第 4 讲的 L 层积木；它自己会把 lecture01 加进 sys.path） # ============================================================================= # 1. softmax 与交叉熵 # ============================================================================= def softmax_naive(Z): """ softmax 的教科书写法：ŷ_c = e^{z_c} / Σ_j e^{z_j}，对每一列（每个样本）分别算。 【思路】先把每个 logit 取指数（变成正数），再除以这一列的总和（变成加起来等于 1 的概率）。 【参数】Z：(C, m) logits 【返回】(C, m)，每一列加起来等于 1 【边界】**这个写法会出问题**：float64 里 e^z 在 z > 709.78 时溢出成 inf，inf / inf = nan； z < −745 时 e^z 下溢成 0，0 / 0 = nan。代码 03 演示 1 专门展示这一点。 调用者需要自己用 np.errstate 关掉警告（否则 python3 -W error 下会报错）。 【numpy 小知识】np.sum(E, axis=0, keepdims=True)：沿 axis=0（类别那一维）求和，结果形状 (1, m)； keepdims=True 保留这一维，除法时才能正确地广播到 (C, m)。 """ E = np.exp(Z) return E / np.sum(E, axis=0, keepdims=True) def softmax(Z): """ 数值稳定的 softmax：先从每一列减去这一列的最大值，再做 softmax。 【思路】softmax 对"每个 logit 加同一个常数"不变： e^{z_c − M} / Σ_j e^{z_j − M} = (e^{z_c} e^{−M}) / (Σ_j e^{z_j} e^{−M}) = e^{z_c} / Σ_j e^{z_j} 取 M = max_j z_j 后，指数里最大的数是 0（e^0 = 1），不会上溢；分母至少是 1，不会出现 0 / 0。 【参数】Z：(C, m) 【返回】(C, m) 概率 【边界】极小的概率仍可能下溢成 0（比如两个 logit 差 1000），这时概率本身就是 0 —— 但不会出现 nan。 要算 log 概率时不要先 softmax 再 log，用 log_softmax。 【DL 联系】PyTorch 的 torch.softmax 内部也是这样做的。 """ Zs = Z - np.max(Z, axis=0, keepdims=True) E = np.exp(Zs) return E / np.sum(E, axis=0, keepdims=True) def logsumexp(Z): """ log Σ_j e^{z_j} 的稳定写法（"log-sum-exp 技巧"）：M + log Σ_j e^{z_j − M}，M = max_j z_j。 【参数】Z：(C, m) 【返回】(1, m) 【思路】和 softmax 一样，先把最大值提出来：log Σ e^{z_j} = log(e^M Σ e^{z_j − M}) = M + log Σ e^{z_j − M}。 右边求和里至少有一项是 e^0 = 1，所以 log 的参数 ≥ 1，既不会溢出也不会 log 0。 """ M = np.max(Z, axis=0, keepdims=True) return M + np.log(np.sum(np.exp(Z - M), axis=0, keepdims=True)) def log_softmax(Z): """ log ŷ_c = z_c − log Σ_j e^{z_j}，稳定版。【参数】Z：(C, m)【返回】(C, m)，每个元素 ≤ 0。 【DL 联系】PyTorch 的 nn.CrossEntropyLoss = LogSoftmax + NLLLoss，正是先算这个（讲义 §3.6）。 """ return Z - logsumexp(Z) def cross_entropy_naive(Z, Y): """ 多分类交叉熵的"先 softmax、再 log"写法：J = −(1/m) Σ_i log ŷ_正确类。 【思路】先用（稳定的）softmax 算出概率，取出正确类别的概率 ŷ_正确类 = Σ_c y_c ŷ_c，再取 −log。 softmax 这一步已经是稳定写法了，所以这里要暴露的是**另一个**问题：取 log 的时机。 【参数】Z：(C, m) logits；Y：(C, m) one-hot 标签 【返回】float 【边界】正确类别的概率小到下溢成 0（比如它的 logit 比最大的 logit 小 1000）时，log 0 = −inf，结果是 inf。 真实的损失是一个有限的大数（见 cross_entropy_from_logits），梯度也完全正常 —— 是"先算概率再取 log"把信息丢掉了。 调用者要自己用 np.errstate 关掉警告。代码 03 演示 1 用它展示问题。 """ m = Z.shape[1] p_correct = np.sum(Y * softmax(Z), axis=0) return float(-np.sum(np.log(p_correct)) / m) def cross_entropy_from_logits(Z, Y): """ 多分类交叉熵的稳定写法：直接从 logits 算，J = −(1/m) Σ_i Σ_c y_c (z_c − logsumexp(z))。 【参数】Z：(C, m) logits；Y：(C, m) one-hot（也可以是每列加起来为 1 的"软标签"） 【返回】float 【思路】one-hot 时，每个样本的损失就是 logsumexp(z) − z_正确类： "所有类的软最大值"减去"正确类的分数"。正确类的分数越突出，损失越接近 0。 【DL 联系】和 PyTorch 的 nn.CrossEntropyLoss()(logits, y) 数学上相同；后者收的是 (m, C) 的 logits 和 (m,) 的整数类别。 """ m = Z.shape[1] return float(-np.sum(Y * log_softmax(Z)) / m) def softmax_jacobian(a): """ 一个样本的 softmax 雅可比矩阵 J[i, j] = ∂ŷ_i / ∂z_j = ŷ_i (δ_ij − ŷ_j)。 【参数】a：(C,) 一个样本的 softmax 输出 【返回】(C, C) 矩阵 = diag(a) − a aᵀ 【思路】i = j 时：∂ŷ_i/∂z_i = ŷ_i − ŷ_i²（和 sigmoid 的 σ(1 − σ) 同一个形状）； i ≠ j 时：∂ŷ_i/∂z_j = −ŷ_i ŷ_j（把 z_j 调大，别的类的概率都会变小 —— "此消彼长"）。 【numpy 小知识】np.diag(a) 把向量放到对角线上；np.outer(a, a) 是外积 a aᵀ，形状 (C, C)。 """ return np.diag(a) - np.outer(a, a) def softmax_model_forward(X, params, hidden="relu"): """ L 层网络 + softmax 输出层的前向传播：[LINEAR → hidden] × (L − 1) → LINEAR → SOFTMAX。 【思路】隐藏层直接复用第 4 讲的 deep_nn.linear_activation_forward；最后一层先用 "linear" 激活得到 logits Z[L]， 再做 softmax。cache 的格式和 deep_nn 完全一样，所以反向可以继续复用 deep_nn 的积木。 【参数】X：(n_x, m)；params：{"W1", "b1", ..., "WL", "bL"}，W[L] 形状 (C, n[L−1])；hidden：隐藏层激活函数名 【返回】(AL, ZL, caches)：AL 是 (C, m) 概率，ZL 是 (C, m) logits，caches 和 deep_nn.L_model_forward 一样 """ caches = [] A = X L = len(params) // 2 for l in range(1, L + 1): act = hidden if l < L else "linear" A, cache = deep_nn.linear_activation_forward(A, params[f"W{l}"], params[f"b{l}"], act) caches.append(cache) ZL = A # 最后一层用的是恒等激活，所以 A 就是 logits return softmax(ZL), ZL, caches def softmax_model_backward(AL, Y, caches, hidden="relu"): """ softmax 输出 + 多分类交叉熵的反向传播。 【思路】输出层起步只有一行：dZ[L] = Ŷ − Y（讲义 §3.5 推导，和二分类的 A − Y 一模一样）。 之后每一层都是第 4 讲的反向积木：linear_backward 算 dW、db（含 1/m）、dA_prev， linear_activation_backward 先穿过激活函数。 【参数】AL：(C, m) softmax 输出；Y：(C, m) one-hot；caches：softmax_model_forward 的 caches；hidden：隐藏层激活 【返回】grads：{"dW1", "db1", ..., "dWL", "dbL"}，和参数同形状（cost J 对参数的梯度，含 1/m） 【DL 联系】这正是 PyTorch 里 nn.CrossEntropyLoss + loss.backward() 自动做的事（代码 03 演示 3 逐数字对照）。 """ grads = {} L = len(caches) dZL = AL - Y # 逐样本的 dZ[L]（第 4 讲约定：不含 1/m） dA_prev, dW, db = deep_nn.linear_backward(dZL, caches[L - 1][0]) grads[f"dW{L}"], grads[f"db{L}"] = dW, db for l in range(L - 1, 0, -1): dA_prev, dW, db = deep_nn.linear_activation_backward(dA_prev, caches[l - 1], hidden) grads[f"dW{l}"], grads[f"db{l}"] = dW, db return grads # ============================================================================= # 2. Batch Normalization # ============================================================================= def batchnorm_forward(Z, gamma, beta, eps=1e-5): """ Batch Norm 的前向（训练时，用这个 mini-batch 自己的均值和方差）。 【思路】对每一个隐藏单元（Z 的每一行）分别做： μ = (1/m) Σ_i z^(i) 这个 mini-batch 上的均值 σ² = (1/m) Σ_i (z^(i) − μ)² 方差（除以 m，不是 m − 1） ẑ^(i) = (z^(i) − μ) / √(σ² + ε) 标准化：均值 0、方差（约）1 z̃^(i) = γ ẑ^(i) + β 再用可学习的 γ、β 把均值和方差"调"到网络想要的值 ε 防止 σ² = 0 时除以 0（PyTorch 默认 ε = 1e−5）。 【参数】 Z : (n, m) 这一层的线性输出（Ng 的 z[l]），每一列一个样本 gamma : (n, 1) 缩放，beta : (n, 1) 平移（注意：这个 β 和动量的 β 没有关系，只是撞了名字） eps : ε 【返回】(Zt, cache)：Zt 是 (n, m) 的 z̃；cache = (Zhat, gamma, inv_std, mu, var) 供反向和"测试时"使用 【边界】m = 1 时 σ² = 0，ẑ = 0，整层输出都等于 β —— 一个样本无法"按 batch 标准化"。PyTorch 在训练模式下 遇到 batch 大小 1 会直接报错（代码 02 演示 2）。 【numpy 小知识】Z.mean(axis=1, keepdims=True)：对每一行求均值，形状 (n, 1)；Z.var 默认就是除以 m 的（有偏）方差。 """ mu = Z.mean(axis=1, keepdims=True) var = Z.var(axis=1, keepdims=True) # ddof=0：除以 m，和 Ng 的公式、PyTorch 训练时一致 inv_std = 1.0 / np.sqrt(var + eps) Zhat = (Z - mu) * inv_std Zt = gamma * Zhat + beta return Zt, (Zhat, gamma, inv_std, mu, var) def batchnorm_backward(dZt, cache): """ Batch Norm 的反向，紧凑公式（讲义 §2.4 推导）。 【思路】设 G = dZ̃（标量损失对 z̃ 的梯度，形状 (n, m)），对每一行： dγ = Σ_i G_i ẑ_i, dβ = Σ_i G_i dẑ = γ G dz_i = (1 / (m σ̃)) · ( m dẑ_i − Σ_j dẑ_j − ẑ_i Σ_j dẑ_j ẑ_j ), σ̃ = √(σ² + ε) 括号里的三项各有含义：第一项是"如果 μ、σ 不随 z 变"时的梯度；第二项来自 μ 也依赖 z_i； 第三项来自 σ 也依赖 z_i。后两项让 dz 满足 Σ_i dz_i = 0（严格成立）和 Σ_i dz_i ẑ_i ≈ 0 —— Batch Norm 的输出对"整列一起平移"完全不敏感，对"整列一起放大"几乎不敏感，所以梯度在这两个方向上的分量为 0 / 几乎为 0。 （"几乎"是因为 ε：ẑ = (z − μ)/√(σ² + ε)，把 z 放大 k 倍，σ² 变成 k²σ²，但 ε 不变，所以 ẑ 会有一点点变化。） 【参数】dZt：(n, m)；cache：batchnorm_forward 返回的 cache 【返回】(dZ, dgamma, dbeta)，形状 (n, m)、(n, 1)、(n, 1) 【边界】这里的 d 都是"同一个标量对各量的梯度"。如果按第 4 讲的约定（dZ 不含 1/m，dW 含 1/m）把它嵌进网络， 那么 dγ、dβ 也要和 dW、db 一样再除以 m；dZ 的公式是线性的，不受影响。 """ Zhat, gamma, inv_std, _, _ = cache m = Zhat.shape[1] dgamma = np.sum(dZt * Zhat, axis=1, keepdims=True) dbeta = np.sum(dZt, axis=1, keepdims=True) dZhat = dZt * gamma dZ = (inv_std / m) * (m * dZhat - np.sum(dZhat, axis=1, keepdims=True) - Zhat * np.sum(dZhat * Zhat, axis=1, keepdims=True)) return dZ, dgamma, dbeta def batchnorm_backward_staged(dZt, cache, Z, eps=1e-5): """ Batch Norm 的反向，"笨办法"：按前向的计算图一步一步倒着走（每一步只用一次链式法则）。 【思路】前向拆成小步： μ = mean(z); d = z − μ; v = mean(d²); s = √(v + ε); ẑ = d / s; z̃ = γ ẑ + β 倒着走： dẑ = γ dz̃ ds = Σ dẑ · d · (−1 / s²) （ẑ = d / s 对 s 求导） dv = ds / (2 s) （s = √(v + ε)） dd = dẑ / s + dv · 2d / m （d 同时出现在 ẑ 和 v 里，两条路的梯度相加） dμ = −Σ dd （d = z − μ） dz = dd + dμ / m （z 同时出现在 d 和 μ 里） 这就是框架的 autograd 在做的事（讲义 §4.4）：把每个小运算的局部导数串起来。 它和 batchnorm_backward 的紧凑公式应当逐数字一致（代码 02 演示 1 验证）。 【参数】dZt：(n, m)；cache：batchnorm_forward 的 cache；Z：前向的输入 (n, m)；eps：和前向相同的 ε 【返回】(dZ, dgamma, dbeta) """ Zhat, gamma, _, mu, var = cache m = Z.shape[1] d = Z - mu s = np.sqrt(var + eps) dgamma = np.sum(dZt * Zhat, axis=1, keepdims=True) dbeta = np.sum(dZt, axis=1, keepdims=True) dZhat = gamma * dZt ds = np.sum(dZhat * d, axis=1, keepdims=True) * (-1.0 / s ** 2) dv = ds / (2.0 * s) dd = dZhat / s + dv * 2.0 * d / m dmu = -np.sum(dd, axis=1, keepdims=True) dZ = dd + dmu / m return dZ, dgamma, dbeta def batchnorm_inference(Z, gamma, beta, running_mean, running_var, eps=1e-5): """ 测试时的 Batch Norm：不用当前 batch 的统计量，而用训练时累积下来的 running mean / var。 【思路】ẑ = (z − μ_running) / √(σ²_running + ε)，z̃ = γ ẑ + β。μ_running、σ²_running 是固定的数， 所以测试时 BN 就是一个普通的"逐元素仿射变换"，每个样本的输出只取决于它自己 —— 不管 batch 里还有谁。 【参数】Z：(n, m)；gamma、beta、running_mean、running_var：(n, 1) 【返回】(n, m) """ return gamma * (Z - running_mean) / np.sqrt(running_var + eps) + beta def ema_update(running, batch_value, beta=0.9): """ 指数加权平均（exponentially weighted average）的一步：running ← β · running + (1 − β) · batch_value。 【思路】第 8 讲的 EWA：β = 0.9 时大约是"最近 1 / (1 − β) = 10 个 batch 的平均"。 【DL 联系】PyTorch 的 nn.BatchNorm1d(momentum=0.1) 写成 running ← (1 − momentum) · running + momentum · batch， 所以 PyTorch 的 momentum = 0.1 对应 Ng 的 β = 0.9 —— 同一件事，参数的含义正好"反过来"（讲义 §2.5）。 """ return beta * running + (1.0 - beta) * batch_value # ============================================================================= # 自检 # ============================================================================= if __name__ == "__main__": rng = np.random.default_rng(0) Z = rng.normal(3.0, 5.0, (4, 8)) gamma, beta = rng.uniform(0.5, 2.0, (4, 1)), rng.normal(size=(4, 1)) G = rng.normal(size=(4, 8)) Zt, cache = batchnorm_forward(Z, gamma, beta) a = batchnorm_backward(G, cache) b = batchnorm_backward_staged(G, cache, Z) print("BN 紧凑公式 vs 计算图逐步：最大差 " + ", ".join(f"{np.max(np.abs(x - y)):.1e}" for x, y in zip(a, b)) + " （dZ, dγ, dβ；应在 1e−15 左右）") Zl = rng.normal(size=(3, 5)) Yl = np.eye(3)[:, rng.integers(0, 3, 5)] print(f"交叉熵：naive {cross_entropy_naive(Zl, Yl):.12f}，from_logits {cross_entropy_from_logits(Zl, Yl):.12f}")