""" ============================================================================= CS230 · 第 9 讲共享模块：PyTorch 训练循环模板（第 10 讲起的 PyTorch 代码可以直接导入） ============================================================================= 对应讲义：lectures/lecture09_tuning_batchnorm_frameworks.md §4.5（"一个完整的训练循环"） ----------------------------------------------------------------------------- 【这个模块是什么】 ----------------------------------------------------------------------------- 第 1–8 讲我们用 numpy 亲手写了训练的每一步。换成 PyTorch 之后，训练循环的"骨架"还是那四步： 前向 forward → 损失 loss → 反向 backward → 更新 update 只是每一步都换成了框架提供的零件： 数据 Dataset / DataLoader 把 (X, y) 包起来，每个 epoch 打乱、切成 mini-batch（第 8 讲的 mini-batch） 模型 nn.Module（这里的 MLP） 在 __init__ 里声明有哪些层，在 forward 里写前向传播；参数自动登记 损失 nn.CrossEntropyLoss 等 直接收 logits（不要自己先做 softmax，讲义 §3.6） 反向 loss.backward() autograd 沿着前向时记录下的计算图倒着走，算出每个参数的 .grad（讲义 §4.4） 更新 torch.optim.SGD / Adam optimizer.step() 用 .grad 更新参数（第 8 讲从零推导过 Momentum、Adam） 模式 model.train() / model.eval() 决定 BatchNorm、Dropout 用"训练时"还是"测试时"的行为（讲义 §2.5） 这个模块把它们组织成一个**可以反复使用的模板**： setup_torch 固定线程数（让结果可复现） ArrayDataset 最简单的 Dataset：把内存里的数组包起来 make_loader 造一个 DataLoader（打乱时用自己的随机数生成器，保证可复现） MLP 全连接网络：[Linear → (BatchNorm) → ReLU → (Dropout)] × (L − 1) → Linear，输出 logits train_one_epoch 训练一个 epoch：标准的"zero_grad → forward → loss → backward → step" evaluate 在 dev / test 上评估：model.eval() + torch.no_grad() fit 把上面两个串起来训练若干个 epoch，记录历史，检测发散 predict 批量预测 count_parameters 数可训练的参数 后面几讲（CNN、RNN、Transformer……）换的只是"模型"那一块，其余部分基本不变。 ----------------------------------------------------------------------------- 【形状约定】PyTorch 的习惯：**每一行一个样本** ----------------------------------------------------------------------------- X：(m, n_x)（图像在第 12 讲起是 (m, 通道, 高, 宽)）；y：(m,) 整数类别（分类）或 (m, 1) 实数（回归） nn.Linear(n_in, n_out) 的 weight 形状是 (n_out, n_in) —— 和 Ng 的 W[l] 完全一样；它算的是 Z = X Wᵀ + b。 ----------------------------------------------------------------------------- 【怎样在别的讲次里导入】 ----------------------------------------------------------------------------- import os, sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture09")) from torch_utils import setup_torch, make_loader, MLP, fit, evaluate 【可复现】本讲义的 PyTorch 代码一律在 CPU 上运行，torch.set_num_threads(1)，并在建模型前 torch.manual_seed(种子)； DataLoader 打乱顺序时用一个单独的 torch.Generator。这样同一台机器上重复运行，输出逐字相同。 （GPU 上的某些运算本身就不是逐位确定的，那是另一回事。） 【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【PyTorch 小知识】【Python 小知识】 直接运行本文件做一个自检（在螺旋数据上训练两次，确认结果逐位相同）： python3 torch_utils.py # 约 2 秒 """ import math import torch from torch import nn from torch.utils.data import DataLoader, Dataset # ============================================================================= # 0. 环境 # ============================================================================= def setup_torch(threads=1): """ 固定 PyTorch 在 CPU 上用的线程数。 【思路】多线程时，一些求和运算的"加法顺序"会随线程调度变化，浮点数的结果就可能在最后几位不同； 对这么小的网络，1 个线程也足够快。 【参数】threads：线程数 【PyTorch 小知识】torch.get_num_threads() 可以查看当前的线程数（默认等于 CPU 核数）。 """ torch.set_num_threads(threads) # ============================================================================= # 1. 数据：Dataset 与 DataLoader # ============================================================================= class ArrayDataset(Dataset): """ 最简单的 Dataset：把内存里的 X、y 包起来，按下标取第 i 个样本。 【思路】PyTorch 把"数据从哪里来"和"怎样分 batch"拆成两个对象： · Dataset 只需要回答两个问题：一共有几个样本（__len__）？第 i 个样本是什么（__getitem__）？ 数据可以在内存里、在硬盘上、甚至是现场生成的 —— DataLoader 不关心。 · DataLoader 负责打乱顺序、把若干个样本拼成一个 mini-batch（第 8 讲）。 PyTorch 自带的 torch.utils.data.TensorDataset 做的事和这个类一样，这里自己写一遍是为了看清楚。 【参数】 X : numpy 数组或张量，第 0 维是样本，形状 (m, ...) y : numpy 数组或张量，形状 (m,) 或 (m, k) x_dtype : X 转成的张量类型（默认 float32：深度学习的惯例，第 4 讲 §4.6 讲为什么） y_dtype : y 的类型。分类任务用 torch.long（整数类别，nn.CrossEntropyLoss 要求的类型）；回归用 float32 【边界】X、y 的样本数不一致时立刻报错，而不是训练到一半才发现。 【Python 小知识】类名后面的 (Dataset) 表示"继承"：ArrayDataset 是一种 Dataset，DataLoader 认得它。 __len__、__getitem__ 是 Python 的"特殊方法"：len(ds) 会调用 ds.__len__()，ds[i] 会调用 ds.__getitem__(i)。 """ def __init__(self, X, y, x_dtype=torch.float32, y_dtype=torch.long): self.X = torch.as_tensor(X, dtype=x_dtype) # as_tensor：能共享内存时不复制 self.y = torch.as_tensor(y, dtype=y_dtype) if len(self.X) != len(self.y): raise ValueError(f"X 有 {len(self.X)} 个样本，y 有 {len(self.y)} 个") def __len__(self): return len(self.X) def __getitem__(self, i): return self.X[i], self.y[i] def make_loader(X, y, batch_size=32, shuffle=False, seed=0, x_dtype=torch.float32, y_dtype=torch.long): """ 造一个 DataLoader：每次迭代给出一个 mini-batch (xb, yb)。 【参数】 X, y : 见 ArrayDataset batch_size : mini-batch 大小（最后一个 batch 可能不满） shuffle : 每个 epoch 开始时是否打乱顺序（训练集打乱，dev / test 不必打乱） seed : 打乱顺序用的种子 —— 每个 DataLoader 用自己的 torch.Generator， 这样"改了别处的随机数"不会影响这里的顺序 【返回】torch.utils.data.DataLoader；for xb, yb in loader: ... 每次得到形状 (batch, ...) 和 (batch, ...) 的张量 【边界】num_workers 用默认的 0（在主进程里取数据）：对内存里的小数据最快，也最容易复现。 数据在硬盘上、需要现场解码图片时，才值得开多个 worker。 【DL 联系】每遍历一次 loader 就是一个 epoch；shuffle=True 时每个 epoch 的顺序都不同（第 8 讲）。 """ gen = torch.Generator().manual_seed(seed) if shuffle else None return DataLoader(ArrayDataset(X, y, x_dtype, y_dtype), batch_size=batch_size, shuffle=shuffle, generator=gen) # ============================================================================= # 2. 模型：nn.Module # ============================================================================= class MLP(nn.Module): """ 全连接网络（multi-layer perceptron）：[Linear → (BatchNorm) → ReLU → (Dropout)] × (L − 1) → Linear。 【思路】在 PyTorch 里写一个模型 = 写一个 nn.Module 的子类： · __init__ 里**声明**有哪些层（nn.Linear、nn.BatchNorm1d……）。凡是作为属性挂在 self 上的层， 它们的参数会被自动登记，model.parameters() 能找到它们，optimizer 才知道要更新谁。 · forward 里**写前向传播**。反向传播不用写 —— autograd 会根据 forward 里实际执行的运算自动求导。 这个 MLP 的输出是 **logits**（最后一层的线性输出 Z[L]），没有 softmax： nn.CrossEntropyLoss 自己会做（数值更稳定，讲义 §3.6）；预测类别时 argmax(logits) 和 argmax(softmax) 一样。 【参数】 layer_dims : [n_x, n[1], ..., n[L]]，和第 4 讲的 layer_dims 含义相同；n[L] 是类别数 C（或回归的输出维数） batchnorm : True 时每个隐藏层的 Linear 后面接 BatchNorm1d，放在 ReLU **之前**（Ng 和原论文的默认做法）， 并且这个 Linear 不要偏置 b（BN 会减掉均值，b 是多余的，讲义 §2.3） dropout : 每个隐藏层 ReLU 之后的 dropout 概率（0 表示不用；第 6 讲讲 dropout） init : "default" = PyTorch 的默认初始化；"he" = 第 7 讲的 He 初始化 W ~ N(0, 2 / n[l−1])、b = 0 【边界】layer_dims 至少要有 2 个数；只有 2 个数时没有隐藏层（[n_x, C] 就是 softmax 回归）。 【PyTorch 小知识】 · nn.Sequential(层1, 层2, ...) 把若干层串起来，forward 时依次调用。 · nn.Linear 的默认初始化是 W ~ U(−1/√n_in, 1/√n_in)（方差 1 / (3 n_in)），比 He 初始化小 √6 倍； 网络很深又没有 BN 时，这会让信号一层层变小（代码 04 演示 4）。 · super().__init__() 必须在最前面调用：它让 nn.Module 准备好"登记参数"的机制。 """ def __init__(self, layer_dims, batchnorm=False, dropout=0.0, init="default"): super().__init__() if len(layer_dims) < 2: raise ValueError("layer_dims 至少要有输入维数和输出维数两个数") layers = [] for n_in, n_out in zip(layer_dims[:-2], layer_dims[1:-1]): # 隐藏层 layers.append(nn.Linear(n_in, n_out, bias=not batchnorm)) if batchnorm: layers.append(nn.BatchNorm1d(n_out)) layers.append(nn.ReLU()) if dropout > 0: layers.append(nn.Dropout(dropout)) layers.append(nn.Linear(layer_dims[-2], layer_dims[-1])) # 输出层：只有线性部分，输出 logits self.net = nn.Sequential(*layers) # *layers：把列表"展开"成多个参数 if init == "he": for module in self.net: if isinstance(module, nn.Linear): nn.init.kaiming_normal_(module.weight, nonlinearity="relu") # N(0, 2 / n_in) if module.bias is not None: nn.init.zeros_(module.bias) elif init != "default": raise ValueError(f"unknown init: {init}") def forward(self, x): """x：(batch, n_x) → logits：(batch, n[L])。""" return self.net(x) def count_parameters(model): """ 可训练参数的个数 = 所有 requires_grad 的参数张量的元素个数之和。 【边界】BatchNorm 的 running_mean / running_var 是 "buffer"（缓冲量）不是参数：它们不参与梯度下降，不计入。 【PyTorch 小知识】p.numel() 是张量的元素个数。 """ return sum(p.numel() for p in model.parameters() if p.requires_grad) # ============================================================================= # 3. 训练与评估 # ============================================================================= def _n_correct(logits, yb): """分类任务：这一批里预测对了几个（argmax 等于标签）。回归任务（y 不是整数）返回 0。""" if yb.dtype != torch.long: return 0 return int((logits.argmax(dim=1) == yb).sum().item()) def train_one_epoch(model, loader, loss_fn, optimizer, device="cpu"): """ 训练一个 epoch：对 loader 里的每个 mini-batch 做一次"前向 → 损失 → 反向 → 更新"。 【思路】循环体里的五行就是整门课的训练循环，和第 1 讲 §2.1 的四步一一对应： logits = model(xb) 1. 前向：autograd 一边算、一边记下计算图 loss = loss_fn(logits, yb) 2. 损失：一个标量（这一批的平均） optimizer.zero_grad() 3a. 把上一步留在 .grad 里的梯度清零（不清零会**累加**！讲义 §4.6） loss.backward() 3b. 反向：沿计算图倒着走，把 ∂loss/∂θ 写进每个参数的 .grad optimizer.step() 4. 更新：θ ← θ − α · (用 .grad 算出的步长) 开头的 model.train() 打开"训练模式"：BatchNorm 用这一批的均值 / 方差（并更新 running 统计量），Dropout 随机丢单元。 【参数】model：nn.Module；loader：训练集的 DataLoader；loss_fn：损失函数对象；optimizer：torch.optim 的优化器； device："cpu" / "cuda" / "mps"，数据要搬到和模型同一个设备上 【返回】(mean_loss, acc)：这个 epoch 里各 mini-batch 损失的（按样本数加权的）平均，以及训练准确率。 【边界】这两个数是"边训练边算"的：每个 batch 用的是那一刻的参数，而且处在训练模式（BN 用 batch 统计量、dropout 打开）。 所以它们和"epoch 结束后在整个训练集上 evaluate"的结果不完全一样，一般略差一点。 【PyTorch 小知识】loss.item() 把只有一个元素的张量变成 Python 的 float，而且**不再连着计算图**。 如果写成 total += loss（不加 .item()），total 会一直拖着整张计算图，内存越用越多。 """ model.train() total_loss, total_correct, total_n = 0.0, 0, 0 for xb, yb in loader: xb, yb = xb.to(device), yb.to(device) logits = model(xb) loss = loss_fn(logits, yb) optimizer.zero_grad() loss.backward() optimizer.step() n = len(yb) total_loss += loss.item() * n total_correct += _n_correct(logits, yb) total_n += n return total_loss / total_n, total_correct / total_n @torch.no_grad() def evaluate(model, loader, loss_fn, device="cpu"): """ 在 dev / test 集上评估：平均损失和准确率。 【思路】两样东西缺一不可： · model.eval()：切到"测试模式"。BatchNorm 改用训练时累积的 running mean / var（每个样本的输出不再依赖 同一批里的其他样本），Dropout 关闭。忘了这一行，测试结果会依赖 batch 的组成（代码 02 演示 2）。 · @torch.no_grad()：这个函数里不记录计算图、不算梯度 —— 更快、更省内存。 注意：no_grad **不会**切换 BN / Dropout 的模式，两者管的是不同的事。 【参数】同 train_one_epoch（不需要 optimizer） 【返回】(mean_loss, acc) 【Python 小知识】@torch.no_grad() 写在 def 上面，叫"装饰器"（decorator）：等于把整个函数体包在 with torch.no_grad(): 里面执行。 """ model.eval() total_loss, total_correct, total_n = 0.0, 0, 0 for xb, yb in loader: xb, yb = xb.to(device), yb.to(device) logits = model(xb) n = len(yb) total_loss += loss_fn(logits, yb).item() * n total_correct += _n_correct(logits, yb) total_n += n return total_loss / total_n, total_correct / total_n def fit(model, optimizer, train_loader, dev_loader=None, loss_fn=None, epochs=10, device="cpu", log_every=0, scheduler=None): """ 完整的训练过程：训练 epochs 个 epoch，每个 epoch 结束后在 dev 上评估一次，记录历史。 【思路】 · 每个 epoch：train_one_epoch → （可选）学习率调度 scheduler.step() → evaluate(dev)。 · 发散检测：训练损失变成 inf / nan 就立刻停止，标记 diverged = True。 （第 1 讲 §6.3 的教训：带着 nan 训练下去，准确率会"安静地"停在某个看似合理的数上。） · log_every > 0 时每隔这么多个 epoch 打印一行进度。 【参数】 model, optimizer : 模型和优化器（优化器要在建好模型之后、用 model.parameters() 创建） train_loader : 训练集 DataLoader（shuffle=True） dev_loader : dev 集 DataLoader；None 时不评估 loss_fn : 损失函数；None 时用 nn.CrossEntropyLoss()（多分类） epochs : epoch 数 device : 见 train_one_epoch log_every : 打印间隔（0 表示不打印） scheduler : 可选的学习率调度器（torch.optim.lr_scheduler 里的对象，第 8 讲讲学习率衰减），每个 epoch 调一次 【返回】(history, info) history : dict，键 "epoch"、"train_loss"、"train_acc"、"dev_loss"、"dev_acc"，每个是长度 = 实际 epoch 数的 list （没有 dev 集时 dev_* 是 nan） info : dict，"diverged"（bool）、"epochs_run"（int） 【边界】返回的 model 是**最后一个 epoch** 的参数（没有 keep_best / 早停；需要时可以在循环里加，第 6 讲讲 early stopping）。 """ if loss_fn is None: loss_fn = nn.CrossEntropyLoss() history = {"epoch": [], "train_loss": [], "train_acc": [], "dev_loss": [], "dev_acc": []} diverged = False epoch = 0 for epoch in range(1, epochs + 1): train_loss, train_acc = train_one_epoch(model, train_loader, loss_fn, optimizer, device) if not math.isfinite(train_loss): diverged = True break if scheduler is not None: scheduler.step() if dev_loader is not None: dev_loss, dev_acc = evaluate(model, dev_loader, loss_fn, device) else: dev_loss, dev_acc = float("nan"), float("nan") for key, val in zip(history, (epoch, train_loss, train_acc, dev_loss, dev_acc)): history[key].append(val) if log_every and epoch % log_every == 0: print(f" epoch {epoch:>4} train loss {train_loss:.4f} acc {train_acc:.3f} " f"dev loss {dev_loss:.4f} acc {dev_acc:.3f}") return history, {"diverged": diverged, "epochs_run": epoch} @torch.no_grad() def predict(model, X, batch_size=1024, device="cpu", x_dtype=torch.float32): """ 批量预测：返回 logits（形状 (m, n[L])）。分类时 logits.argmax(dim=1) 就是预测的类别。 【思路】同样要 model.eval() + no_grad。分批算是为了数据很大时不一次性占满内存。 【参数】X：numpy 数组或张量，形状 (m, ...)；其余见上 【返回】CPU 上的张量 (m, n[L])，不连着计算图 【PyTorch 小知识】torch.cat(列表, dim=0) 把若干个张量沿第 0 维拼起来。 """ model.eval() X = torch.as_tensor(X, dtype=x_dtype) outs = [model(X[i:i + batch_size].to(device)).cpu() for i in range(0, len(X), batch_size)] return torch.cat(outs, dim=0) # ============================================================================= # 自检：在螺旋数据上训练两次，确认（1）能学会，（2）同一个种子的结果逐位相同 # ============================================================================= def _selftest_once(seed): import os import sys import numpy as np sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from multiclass_data import make_spirals X, y = make_spirals(200, 3, np.random.default_rng(1)) Xd, yd = make_spirals(200, 3, np.random.default_rng(2)) torch.manual_seed(seed) # 模型的初始化用全局随机数，所以建模型之前设种子 model = MLP([2, 64, 64, 3]) opt = torch.optim.Adam(model.parameters(), lr=1e-2) hist, info = fit(model, opt, make_loader(X, y, 64, shuffle=True, seed=seed), make_loader(Xd, yd, 256), epochs=60) return hist, info if __name__ == "__main__": setup_torch(1) h1, i1 = _selftest_once(0) h2, _ = _selftest_once(0) print(f"螺旋数据（3 类），MLP [2, 64, 64, 3]，Adam lr 0.01，60 个 epoch：") print(f" 最后一个 epoch：train loss {h1['train_loss'][-1]:.4f}，dev 准确率 {h1['dev_acc'][-1]:.3f}，" f"发散：{i1['diverged']}") same = all(h1[k] == h2[k] for k in h1) print(f" 同一个种子再训练一次，历史记录逐位相同：{same}") print(f" 可训练参数 {count_parameters(MLP([2, 64, 64, 3]))} 个；加 BatchNorm 后 " f"{count_parameters(MLP([2, 64, 64, 3], batchnorm=True))} 个（少了 2 个隐藏层的偏置 64 + 64，多了 γ、β 共 4 × 64）")