PyTorch 速成¶
八章把 PyTorch 用熟:张量与形状、索引与掩码、autograd、nn.Module、数据与训练循环、保存与复现、调试与提速。这本书只讲怎么用——每一节都是一段几秒钟就能跑完的代码加上它真实的输出,读完你能看懂也能写出主流框架里的训练代码。想知道这些东西怎么实现,去 CUDA 手册的「框架与编译器」几章;想把它们用起来完整训练一个模型,去《从零训练一个小模型》。
8 章5 道练习题每段代码都在 CPU 版 PyTorch 上真跑过
这本书适合谁¶
- 会 Python,看得懂神经网络在算什么,但 PyTorch 只停留在"抄得动别人的训练脚本";
- 读源码时经常卡在形状、
gather、detach、state_dict这类地方; - 想把"形状对不上""梯度是 None""显存莫名其妙涨"这类问题自己定位掉。
读完并跟着敲完,你应该能:
- 一眼看出一段张量操作的形状怎么变,知道什么时候是视图、什么时候复制了数据;
- 用
gather/scatter_/masked_fill/einsum写出注意力和损失函数里的那些操作; - 讲清
no_grad、detach、eval()各自管什么,为什么它们不能互相替代; - 写出一个带评估、调度和裁剪的训练循环,并且中断之后能原样续上;
- 读懂报错,用 profiler 找到最慢的算子。
八章的路线¶
| 章节 | 做什么 | 学完能做什么 | 建议用时 |
|---|---|---|---|
| (一)张量 | 形状、dtype、设备,以及哪些操作共享内存 | 不再被 dtype 和设备的报错卡住 | 1 小时 |
| (二)形状的功夫 | view / reshape / permute、广播、einsum | 看一眼就知道形状怎么变,能写多头注意力的下标 | 2 小时 |
| (三)索引、归约与掩码 | 视图与拷贝、gather / scatter_、因果掩码、dim 与 keepdim | 写得出交叉熵和注意力里的取数操作 | 2 小时 |
| (四)autograd 怎么用 | backward、梯度累加、no_grad 与 detach | 讲清梯度为什么是 None,知道哪些原地操作会炸 | 2 小时 |
| (五)nn.Module | 参数与 buffer、state_dict、train / eval | 写出结构清晰、存得下也加载得回来的模型 | 2 小时 |
| (六)数据与训练循环 | Dataset、DataLoader、collate_fn,一个完整的训练 | 独立写出一个能跑的训练脚本 | 3 小时 |
| (七)保存、加载与复现 | checkpoint 存什么、种子与确定性 | 中断之后续训的结果和没中断一样 | 1 小时 |
| (八)调试与提速 | 读懂报错、经典的坑、inference_mode、autocast、profiler | 自己定位形状、梯度和性能问题 | 2 小时 |
八章彼此独立,哪一章卡住就单独看哪一章;代码都在 CPU 上几秒钟跑完。
接下来往哪走¶
| 想继续的方向 | 去哪本 |
|---|---|
| 把这些用起来,完整训练一个会写《三国演义》的小模型 | 从零训练一个小模型 |
| 张量在内存里到底长什么样、autograd 怎么建图、torch.compile 做了什么 | CUDA 进阶手册的「框架与编译器」 |
| Transformer 的每个部件为什么这样设计 | 大模型原理 |
| 一张卡放不下的时候:DDP、ZeRO、张量并行 | 分布式训练手册 |
| 写得地道一点的 Python | Python 进阶 |
怎么验证的¶
- 每一页的代码都在 CPU 版 PyTorch 上实际运行,页面上的输出与运行结果逐行一致;
- 报错信息也是真跑出来的——这本书里出现的每一条
RuntimeError都可以自己复现。
# 需要 CPU 版 PyTorch(2.4 以上)与 numpy
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install numpy
发现错误或有更好的写法?在 GitHub 上编辑此页提交问题