6.1 框架与算力：训练流水线的基础设施
6.1 框架与算力：训练流水线的基础设施
本节摘要：框架负责把前八章的手工账自动化（自动微分、GPU 调度、算子融合），算力决定练得多快、多便宜。本节给出 PyTorch 与 TensorFlow 的选型判断、GPU 之所以快的两笔算术、显存与乘加的预算方法，以及量化与混合精度这两笔最常用的省钱账。
从演示到交付的分水岭
笔记本上跑通 TinyCNN 与给业务方交付一个模型，中间隔的是基础设施。框架层面，全世界的深度学习代码基本收敛在两大生态：PyTorch 以动态图和 Python 原生体验成为研究与新项目的默认，TensorFlow 配套的部署工具链（面向手机与嵌入端的转换器）在端侧交付上历史更久。选择的判断其实简单：新项目、要发论文、团队用 Python——PyTorch；存量系统、强端侧约束——认真评估 TensorFlow 系的部署链。与其押宝框架，不如把通用的账算会：显存、乘加、精度，这三本账换哪个框架都成立。
学习目标
阅读完本节，你应当能够：
- 心算任意批次的输入显存占用（第一章公式的高频复用）；
- 复算一段卷积网络的乘加量，给"训练要多久"一个数量级判断；
- 解释 GPU 快的两笔算术：大规模并行与高带宽显存；
- 说出混合精度与量化各自的省钱原理与精度风险。
一、两笔基础账：显存与乘加
显存的第一笔账在 1.3 节立过规矩：元素个数乘四字节（float32）。批 64 的 224 分辨率输入约 36.8 MiB——这只是入场费，训练时中间激活与反向梯度才是大头，粗略估计是输入的几十倍，这也是"批大小一加就溢出"的算术根源。乘加账在 2.6 节练过：AlexNet 的卷积段每图约 1.08 GMACs，六千二百万参数的它处理一张图就要十亿次量级乘加，训练集一百万张、跑几十个轮次，总量直接跳到 10 的 18 次方量级——没有 GPU，这笔账靠 CPU 要按月计。
import torch print("CUDA 可用:", torch.cuda.is_available()) if torch.cuda.is_available(): print("设备名:", torch.cuda.get_device_name(0)) x = torch.randn(64, 3, 224, 224) print("批次输入显存:", round(x.numel() * 4 / 1024**2, 1), "MiB") # 36.8 MiB
GPU 快的算术不神秘。其一，吞吐：训练的主体是大矩阵乘法，天然可拆成海量独立的小乘加，GPU 上千个核心同时开工，CPU 十几个核心只能望洋兴叹；其二，带宽：中间特征图在显存与计算单元之间搬运频繁，GPU 显存带宽是系统内存的数倍，搬运不吃亏。两者叠加，矩阵越大 GPU 优势越大——这就是为什么小模型调试在 CPU 上无妨，正式训练必上 GPU。
def conv_macs(cin, cout, k, h, w): return cout * h * w * cin * k * k # 2.1 节的乘加公式 macs = 0 macs += conv_macs(3, 96, 11, 55, 55) # AlexNet 卷积段逐层累计 macs += conv_macs(96, 256, 5, 27, 27) macs += conv_macs(256, 384, 3, 13, 13) macs += conv_macs(384, 384, 3, 13, 13) macs += conv_macs(384, 256, 3, 13, 13) print(f"卷积段乘加合计: {macs/1e9:.2f} GMACs/图") # 1.08 GMACs
二、省钱的两种姿势：混合精度与量化
精度每降一位，显存与带宽省一半。混合精度训练把大多数算子换成半精度（float16 或 bfloat16），权重主副本与累加保留单精度防漂移，框架的自动包装一层搞定——速度通常提升近倍、显存近乎减半，代价是少数数值敏感的层（归一化、损失）需要留在高精度，现代实现已处理成熟。量化是部署端的省钱术：把训练好的 float32 权重压到 int8（理论上体积缩到四分之一），端侧推理快、省电；代价是精度回落，需要用校准数据把激活范围量准，或做量化感知训练把回落压到最小。
scaler = torch.amp.GradScaler() # 混合精度的梯度缩放器 # 训练循环只需三处改动（伪代码骨架）： # with torch.autocast(device_type="cuda"): # 前向自动换半精度 # loss = crit(net(x), y) # scaler.scale(loss).backward() # 缩放梯度防下溢 # scaler.step(opt); scaler.update()
部署链的最后一段通常是：训练框架导出中间表示（如 ONNX），推理引擎（TensorRT 一类）按目标硬件编译优化，端侧再叠加量化。这条链上每一环都是"先对账、再换格式"——乘加多少、带宽多少、精度掉多少，前三章的公式一路通用。算力预算的完整心法一句话：训练看显存与乘加，部署看时延与体积，两本账都要先算后买。
追问两则
问：梯度累积怎么算等效批？答：累积 k 次小批再更新一次，等效批大小 = 小批 × k——显存只装得下批 8 时，累积四次就近似批 32 的梯度质量，代价是每轮时间不变而更新次数变少，学习率一般不用改。问：什么任务 CPU 训练也够？答：模型十万参数以下、图很小的场景（教学网络、嵌入式原型）CPU 几分钟能出一轮；但只要模型上了百万参数或图上了 64 分辨率，GPU 的账就开始碾压，别在错误的地基上耗时间。
巡检记录
- 框架判断：研究与新项目默认 PyTorch，端侧强约束认真评估 TensorFlow 部署链；账本比站队重要；
- 两笔基础账：显存 = 元素数乘 4 字节（输入 36.8 MiB 起步、激活是大头），乘加用 2.1 节公式逐层累计（AlexNet 卷积段 1.08 GMACs/图）；
- GPU 两算术：千核并行吃矩阵乘，高带宽喂特征搬运——负载越矩阵化优势越大；
- 省钱姿势：混合精度训练近倍提速、显存减半；int8 量化部署体积缩四倍，精度用校准找回。
⚠️ 常见坑：显存溢出第一反应换更大显卡。先减批大小、开混合精度、加梯度累积（小批多次等效大批）——三招通常能把预算压下来一半，账算清之前别急着花钱。
工具备齐，该审视风险了：这套系统会被怎么骗、它的判断能不能向人交代清楚。下一节走进可解释性与对抗样本。