构建者指南 — 模块、参数与 GPU
这一章讲两件事: 让复杂网络可以「搭积木」的 module 抽象, 和让代码在 GPU 上真正跑起来的设备模型。 这是全书从「模型原理」转向「工程习惯」的一章,篇幅最短, 但没有它,后面 ResNet-152 这种上百层的网络根本没法下手。
1. module:层、块、整个模型,是同一种东西
回看我们已经见过的对象,会发现一个递归(自己套自己、同种结构一层包一层)结构1:
- 一个神经元:收输入 → 出标量 → 有自己的参数;
- 一层:收输入 → 出向量 → 有自己的参数;
- 整个模型:收特征 → 出预测 → 参数 = 所有层的参数之和。
三者是同构的。 那么比层大、比模型小的东西呢? ResNet-152 有 152 层,但不是 152 个各不相同的设计—— 它们是少数几种「层组」重复堆出来的2。 一层一层手写这种网络,既枯燥又易错。
框架给出的抽象叫 module(模块):
一个类,必须定义 forward(输入怎么变输出),
参数由框架登记管理,反向传播由自动微分白送3。
模块的关键性质是可以递归组合:
层是模块,几层拼成的块是模块,整个模型还是模块——
nn.Sequential 就是最简单的那种:一张有序清单,
前一个模块的输出接后一个的输入。
调用 net(X) 时会转到 forward——net(X) 这种「像调函数一样调对象」的写法,
是 Python 在语法(一门语言规定的写法规则)上给的甜头,行话叫「语法糖」。
你写模型只需要填两样东西:构造函数里建哪些模块,forward 里怎么连。
「定义结构」与「写前向」分离,是这套抽象的全部精髓。
2. 参数:访问、共享与「先欠着」
模块的参数不只是「存在某处的矩阵」,有三个工程问题要答。
访问与初始化。 框架给每个参数登记在册(net[i].weight、net[i].bias),
初始化方案可以按层指定(第 06 章的 Xavier 是常用默认)。
共享参数。 有时候希望两处用同一份参数 (比如后面 seq2seq 里编码器和解码器共享 embedding)。 做法不是复制数值,而是让两个位置引用同一个参数对象—— 梯度在反传时自然汇总到这一份上。
延迟(推迟到框架第一次见到真实输入形状时才进行的)初始化(lazy init)。 手写全连接层要声明输入维度(比如 784), 换个数据集就得改。框架可以第一次见到真实输入时才分配参数4: 形状到那时才确定,代码里不用写死任何维度。 这不是语法糖,是把「形状推断」这件易错的事外包给了数据本身。
剩下两件小事实用而直白:自定义层(不带参数的变换,比如把输入居中化;
或带参数的完整新模块);以及把参数存盘/读盘(state_dict)——
模型 = 结构代码 + 参数文件,两者分开存。
主走查:一个 nn.Sequential,从 (2, 20) 走到 (2, 10)
把 §1、§2 的规矩压进原书的演示网:
nn.Sequential(nn.LazyLinear(256), nn.ReLU(), nn.LazyLinear(10)),
喂一批 X——2 条样本、每条 20 维5:
X:(2, 20) ← 批大小 2,特征 20
│ net[0] LazyLinear(256):W₀(20 维进、256 维出)+ b₀(256)
▼
(2, 256)
│ net[1] ReLU:无参数,逐元素 max(0, ·)
▼
(2, 256)
│ net[2] LazyLinear(10):W₂(256 维进、10 维出)+ b₂(10)
▼
(2, 10) ← 每条样本 10 个分数
走完盯三件事。
参数归属。 W₀/b₀ 挂在 net[0] 名下,W₂/b₂ 挂在 net[2] 名下,
ReLU 一无所有——按层索引(按编号点名取用)就能取出(net[2].weight、net[2].bias)6。
「模型的参数 = 各 module 登记的参数之和」,递归在这条网上兑现。
形状何时确定。 建网那一刻框架不知道输入是 20 维, 第一层连参数都还没分配7;X 一到,把 20 代入, 第一层形状落定,再顺计算图一层层推下去, 全部形状齐了才统一初始化8——这就是 §2 说的「先欠着」在运行时的样子。
共享参数。 若让清单里两个位置放同一个 shared 层
(比如第 2、4 个模块是同一个 LazyLinear(8) 对象),
前向照常一层层走,但参数只有一份——
框架保证那是「同一个对象」而不是「同值的两份」,
反传时两处的梯度都汇总进这一份9。
3. GPU:设备模型的三条铁律
第 01 章说了深度学习靠 GPU 起飞。这一节是用 GPU 的规矩。 先补一个数字:GPU 算力自 2000 年以来大约每十年涨 1000 倍10。
铁律一:运算双方必须在同一设备上。 想对两个张量做运算,它们必须都待在 CPU 内存、或同一张 GPU 的显存里。 把 CPU 上的 X 和 GPU 上的 Y 直接相加,框架报错11。
为什么宁可报错也不悄悄帮你搬?因为跨设备复制远慢于计算—— 如果框架默默搬了,你根本意识不到自己写出了慢代码。 作者的原话是:我们要你 100% 确定想做这件慢事,才让你做12。
铁律二:传输次数比传输量更要命。 许多次小传输远不如一次大传输——每次传输都有固定开销, 而且会造成设备间互相等待。 「像排队买咖啡,不如提前电话下单,到了就取」13。
铁律三:日志也会卡住整条流水线。
一个高频踩坑:每个 minibatch 算完 loss 就 .item() 搬回 CPU 打印——
这一搬不仅慢,还会撞上 Python 的全局解释器锁(GIL),
让所有 GPU 停下来等 Python14。
正确姿势:在 GPU 上记账,攒一大批再一次性搬回来。
把模型放上 GPU 只有一件事:模型参数和输入数据在同一个设备, 训练就和之前一模一样——设备模型对算法是透明的,对性能不是。
4. 作者的判断与证据
书里给了证据的: module 的递归结构(框架源码层面); 设备间复制的相对成本(「much slower than computation」); GIL 卡 GPU 的具体场景(每批打印 loss)。
经验之谈(作者未给证明,属实践共识): 「多次小传输不如一次大传输」; 「lazy init 省掉的形状错误比它带来的隐式性更值」。
5. 边界与局限
- 本章只覆盖单 GPU 的设备模型;多 GPU 与分布式是第 15 章;
- module 抽象各框架细节不同(PyTorch 的
nn.Module、Flax 的 dataclass 风格), 拆解讲的是共同骨架; - 参数存取的安全与版本问题(加载旧代码存的参数)原书未展开。
6. 可带走的
- 层、块、模型是同一种对象(module);写模型 = 声明模块 + 写 forward;
- 共享参数靠「引用同一对象」,梯度自动汇总;
- 延迟初始化:形状让数据说话;
- 模型 = 结构代码 + 参数文件,分开存;
- 运算双方必须同设备;框架报错是在救你,不是在烦你;
- 多次小传输不如一次大传输;
- 别把 loss 逐批搬回 CPU 打印——GIL 会停住所有 GPU。
7. 原文地图
| 主题 | 原书章 | 原文位置 |
|---|---|---|
| 层组与 module 抽象 | Layers and Modules | text/33-layers-and-modules.txt:55(搜「groups of layers」) · text/33-layers-and-modules.txt:68(搜「module」) |
| 递归组合、只写 forward | Layers and Modules | text/33-layers-and-modules.txt:74(搜「recursively」) · text/33-layers-and-modules.txt:84(搜「forward propagation method」) |
| 主走查:演示网与 (2, 20) 输入 | Layers and Modules | text/33-layers-and-modules.txt:147(搜「LazyLinear(256)」) · text/33-layers-and-modules.txt:149(搜「rand(2, 20)」) |
| 参数按层访问 | Parameter Management | text/34-parameter-management.txt:139(搜「state_dict」) |
| 共享同一参数对象 | Parameter Management | text/34-parameter-management.txt:284(搜「shared = nn.LazyLinear(8)」) · text/34-parameter-management.txt:294(搜「same object rather than」) |
| 延迟初始化 | Lazy Initialization | text/36-lazy-initialization.txt:25(搜「defers initialization」) |
| 形状逐层推定 | Lazy Initialization | text/36-lazy-initialization.txt:196(搜「plugging in the value of 20」) · text/36-lazy-initialization.txt:204(搜「Once all parameter shapes are known」) |
| 每十年 1000 倍 | GPUs | text/39-gpus.txt:12(搜「factor of 1000」) |
| 同设备铁律 | GPUs | text/39-gpus.txt:274(搜「same device」) · text/39-gpus.txt:354(搜「cannot find data」) |
| 复制远慢于计算、100% 确定 | GPUs | text/39-gpus.txt:458(搜「much slower than computation」) · text/39-gpus.txt:459(搜「100% certain」) |
| 小传输 vs 大传输、咖啡 | GPUs | text/39-gpus.txt:469(搜「many small operations」) |
| GIL 卡 GPU | GPUs | text/39-gpus.txt:484(搜「global interpreter lock」) |