实战演练:手写 conv2d 并与 PyTorch 对齐
案例篇把 5×5 的 81 次乘法算到了每一位。真实网络里,一层 3×3 卷积就是 22 万次乘加,一个 MNIST 小 CNN 每张图 103 万次。这一篇把「算一次」变成「算一亿次」:先写 im2col,再写朴素四重循环,再和 F.conv2d 逐元素对齐,最后把卷积包成可训练层,量出慢在哪。
所有代码在 assets/run_experiments.py,所有数字由它写入 assets/metrics.json。从头跑一遍:
cd "第 4 章:CNN 原理/assets"
/private/tmp/dnn-venv/bin/python run_experiments.py本机 29.9 秒结束,输出 12 组对齐、4 项梯度、3 档速度基准。

<!-- 配图:五张卡片从左到右排列:im2col、朴素循环、12 组对齐、自定义层、速度基准;每张卡片底部给出产物或关键数字,最后一张 98.4× -->
零、先把「对」写死
「和 PyTorch 对齐」不是一句感觉,而是三条可执行的判据。先写判据,再写代码,后面所有优化才有裁判。

<!-- 配图:三张卡片分别给出数值 1e-5、梯度 1e-10、速度 98.4×;底部环境条列出 Python 3.9.6、torch 2.8.0、numpy 2.0.2、4 线程 -->
| 项目 | 判据 | 实测 |
|---|---|---|
| 数值(float64) | 12 组输出逐元素最大差 < 1e-5 | 9.24e-14 |
| 数值(float32,附加) | 同上,仅作压力测试 | 5.34e-5,9/12 达标 |
| 梯度 | 前向、grad_x、grad_w、grad_b 最大差 < 1e-10 | 1.24e-14 |
| 速度 | 能解释朴素循环为什么慢 | 0.783 s → 0.00796 s,98.4× |
环境是固定的:Python 3.9.6、torch 2.8.0、numpy 2.0.2、torch.set_num_threads(4)、随机种子 20261031。环境一变,后文的毫秒数不能直接横比,但误差量级和「快慢关系」不变。
这里有一个容易含糊的地方:「误差小于 1e-5」必须写清 dtype。同样的实现,float64 的最大差是 9.24e-14,float32 是 5.34e-5;前者轻松过线,后者有三组过不去。这不是实现错,后面 §3.5 会拆开。
一、im2col:把每个窗口摊成一列
1.1 从 9 个数到一列
卷积的输出位置彼此独立,每个位置只看一个 C×K×K 的小块。把这些小块按「输出位置」排成列,就把卷积变成了矩阵乘。

<!-- 配图:左边 4×4 输入高亮左上 3×3 窗口;右边 9×4 矩阵,第一列就是该窗口的 9 个数,其余三列对应另外三个窗口 -->
4×4 输入、3×3 核、stride 1,输出是 2×2,一共 4 个位置。把每个位置的 9 个输入按 (通道, 核行, 核列) 的顺序排好,就得到 9 行 4 列:
P1 = [1, 2, 3, 5, 6, 7, 9, 10, 11]
P2 = [2, 3, 4, 6, 7, 8, 10, 11, 12]
P3 = [5, 6, 7, 9, 10, 11, 13, 14, 15]
P4 = [6, 7, 8, 10, 11, 12, 14, 15, 16]列高是 C_in×K×K,列数是 N×H_out×W_out。真实大案例里 C_in=16、K=3,列高 144;64×64 输出、每张图 4,096 列;8 张图拼在一起是 32,768 列。
1.2 展开只有三层循环
im2col 不需要四层循环。它只遍历「通道、核行、核列」三层,每一层把一个固定偏移的采样面切片、拉平、塞进一行:
# 第 4 章:CNN 原理/assets/run_experiments.py :: im2col()
def im2col(x, k, stride=1, padding=0, dilation=1):
"""x: (N,C,H,W) float32/64。按 C×K×K 行、输出位置列排布。"""
N, C, H, W = x.shape
if padding:
xp = np.pad(x, ((0, 0), (0, 0), (padding, padding), (padding, padding)))
else:
xp = x
KH = (k - 1) * dilation + 1
OH = (H + 2 * padding - KH) // stride + 1
OW = (W + 2 * padding - KH) // stride + 1
cols = np.empty((N, C * k * k, OH * OW), dtype=x.dtype)
col = 0
for c in range(C):
for ki in range(k):
for kj in range(k):
patch = xp[:, c,
ki * dilation:ki * dilation + OH * stride:stride,
kj * dilation:kj * dilation + OW * stride:stride]
cols[:, col, :] = patch.reshape(N, -1)
col += 1
return cols, OH, OW三个循环变量各管一件事:
c选输入通道,决定这一行来自哪张特征图;ki选核内第几行,起点是ki*dilation,所以 dilation 在这里生效;kj选核内第几列,终点是kj*dilation + OW*stride。
最关键的是切片里的 :stride:同一个核内偏移,在输出位置之间是按 stride 跳的,不是按 dilation 跳。dilation 只改「起点」,stride 只改「步长」,两件事不能混。案例篇 §3 的 7×7、d=2 例子,就是靠这个切片取到行/列 0、2、4。
以 4×4 输入、K=3、stride 1 为例,cols 的 9 行是按 (c, ki, kj) 的顺序填出来的,第一列对应窗口 P1:
| 行号 | (ki, kj) | P1 中取的格子 | 值 |
|---|---|---|---|
| 0 | (0, 0) | (0, 0) | 1 |
| 1 | (0, 1) | (0, 1) | 2 |
| 2 | (0, 2) | (0, 2) | 3 |
| 3 | (1, 0) | (1, 0) | 5 |
| 4 | (1, 1) | (1, 1) | 6 |
| 5 | (1, 2) | (1, 2) | 7 |
| 6 | (2, 0) | (2, 0) | 9 |
| 7 | (2, 1) | (2, 1) | 10 |
| 8 | (2, 2) | (2, 2) | 11 |
行号是 col 的递增顺序,列号是输出位置。c 每变一次,行号跨过 K×K,进入下一个输入通道;ki 每变一次,行号跨过 K,进入下一行核。这个顺序必须和权重 w.reshape(O, -1) 完全一致,否则矩阵乘会把「通道 1 的第 2 行」和「通道 2 的第 1 行」乘在一起,形状对、数值全错。
1.3 矩阵乘:维度必须对上
展开之后,权重只需要按同样的顺序摊成 O×CKK:

<!-- 配图:W 为 4×9、cols 为 9×6、out 为 4×6;乘号与等号连接三个矩阵,黄色格标出 out(0,0) 由 W 第 0 行与 cols 第 0 列点积得到 -->
W: O × CKK 4 × 9
cols: CKK × L 9 × 6
out: O × L 4 × 6CKK 必须同时是 W 的列数和 cols 的行数。真实大案例是 32×144 · 144×32,768 = 32×32,768:一次矩阵乘,把 1,048,576 个输出位置的 144 次乘加全部交给 BLAS。
代码只有一行核心:
# 第 4 章:CNN 原理/assets/run_experiments.py :: im2col_conv2d()
def im2col_conv2d(x, w, b, stride=1, padding=0, dilation=1):
cols, OH, OW = im2col(x, w.shape[2], stride, padding, dilation)
N = x.shape[0]
O = w.shape[0]
# einsum 直接表达 (N,CKK,L)×(O,CKK)→(N,O,L)。
# 本机 NumPy 2.0.2 的 matmul 在上述部分形状上会误报 divide by zero。
out = np.einsum("ok,nkl->nol", w.reshape(O, -1), cols, optimize=True)
out = out.reshape(N, O, OH, OW)
out += b.reshape(1, O, 1, 1)
return out踩坑记录(一):np.matmul 的误报
- 现象:把
einsum换成np.matmul(w.reshape(O, -1), cols),在1×1×1×1与几个大形状上抛divide by zero encountered/overflow,但同一组数据用einsum结果正常。 - 原因:这是本机 NumPy 2.0.2 在部分形状组合上的 BLAS 路径问题,不是输入里有零或溢出;换成
einsum(optimize=True)走另一条 contraction 路径就消失。 - 修复:保留
einsum,并在注释里写清原因。不要为了「看起来更像矩阵乘」而改回matmul,除非你在目标环境重新验过。
这类坑的教训不是「matmul 不能用」,而是:依赖库的数值路径不是数学定义,换 API 要重跑对齐。
1.4 四个阶段
把上面的步骤串起来,一次 im2col 卷积只有四个阶段:

<!-- 配图:输入 → im2col → matmul → reshape(+bias) → 输出五个节点,im2col 与 matmul 用蓝框;底部三条说明顺序不能换 -->
输入 8×16×64×64
→ padding(补 1 圈零,变 8×16×66×66)
→ im2col(144×32,768 的列矩阵)
→ matmul(32×144 · 144×32,768 → 32×32,768)
→ reshape + bias(8×32×64×64)顺序不能换。先补零再展开,列里才会包含边界零;先矩阵乘再 reshape,通道与空间位置才不会错位。任何一步顺序错了,形状可能仍然「看起来对」,但数值会整体错位。
1.5 代价:9 倍显存
im2col 不是免费的。它把输入复制了 9 倍:

<!-- 配图:左边一张 2 MB 输入方块,右边 3×3 共 9 张同样大小的方块代表 18 MB 列矩阵;底部三张卡片给出 2.0 MB、18.0 MB、9.0× -->
大案例的账:
输入: 8×16×64×64×4 B = 2,097,152 B = 2.0 MB
列矩阵:32,768×144×4 B = 18,874,368 B = 18.0 MB
膨胀: 18.0 / 2.0 = 9.0×9 倍不是巧合:K×K=9,每个输入像素会被 9 个窗口复用,展开时就复制 9 份。显存不够时,标准做法是把输出位置 L 切块,每次只展开一部分,算完就丢。
设计取舍:朴素循环 vs im2col
| 方案 | 收益 | 代价 |
|---|---|---|
| A:朴素四重循环 | 零额外显存;代码短,和公式一一对应;改 dilation/stride 只改切片 | 大案例 0.783 s;1,048,576 次 Python 调度 |
| B:im2col + 矩阵乘 | 大案例 0.00796 s,快 98.4×;矩阵乘交给 BLAS | 18 MB 列矩阵,9 倍膨胀;多一步展开与 reshape |
这一篇两个都写:A 用来做正确性基线,B 用来跑速度。真实框架不会让你二选一,它把 B 的展开和矩阵乘融进内核(implicit GEMM),连列矩阵都不物化;但那是 C++/CUDA 的事,Python 侧先看清楚这两条路。
1.6 用 F.unfold 交叉验证
自己写的 im2col 最容易错在行顺序。最快验证方法是和 PyTorch 的 F.unfold 比:它返回的就是 (N, C×K×K, L) 的列矩阵,语义与我们的 cols 完全一致。
# 第 4 章:CNN 原理/assets/run_experiments.py :: im2col() 的对照实验(临时脚本)
x = np.random.default_rng(1).standard_normal((2, 3, 7, 7)).astype(np.float32)
cols, OH, OW = im2col(x, 3, stride=2, padding=1, dilation=2)
cols_t = F.unfold(torch.from_numpy(x), (3, 3),
dilation=2, padding=1, stride=2)
print(cols.shape, tuple(cols_t.shape)) # 两边都是 (2, 27, 9)三组配置(5×5 单通道、7×7 双通道 + d=2、4×4 双通道 + K=2)逐元素比较,结果完全一致。这一步值得单独跑,因为它把「切片写错」和「矩阵乘写错」分开了:如果 cols 和 F.unfold 对不上,后面所有对齐都是在错的基础上调。
F.unfold 的输入必须是 float 张量,np.float64 要先 torch.from_numpy 再转;它不接受整数张量。这个限制和第 4.2 节的自定义层一致:卷积的输入、权重、bias 都按 float 处理。
二、朴素四重循环:正确性基线
2.1 四层循环
先写最笨的版本。它不追求快,只追求「每一笔都能和公式对上」:

<!-- 配图:四层嵌套框 N→O→行→列,最内层黄色区域给出 patch 切片与 np.dot 一行;底部总计 1,048,576 次最内层、150,994,944 次乘加 -->
# 第 4 章:CNN 原理/assets/run_experiments.py :: naive_conv2d()
def naive_conv2d(x, w, b, stride=1, padding=0, dilation=1):
"""朴素四重循环:N、O、输出行、输出列;核内用 numpy 点积。"""
N, C, H, W = x.shape
O, _, K, _ = w.shape
if padding:
xp = np.pad(x, ((0, 0), (0, 0), (padding, padding), (padding, padding)))
else:
xp = x
KH = (K - 1) * dilation + 1
OH = (H + 2 * padding - KH) // stride + 1
OW = (W + 2 * padding - KH) // stride + 1
wf = w.reshape(O, C * K * K)
out = np.zeros((N, O, OH, OW), dtype=x.dtype)
for n in range(N):
for o in range(O):
for i in range(OH):
for j in range(OW):
patch = xp[n, :,
i * stride:i * stride + KH:dilation,
j * stride:j * stride + KH:dilation]
out[n, o, i, j] = np.dot(patch.reshape(-1), wf[o]) + b[o]
return out循环顺序是 N → O → 输出行 → 输出列。最内层不再写 C 和 K 两层,而是把整个 C×K×K 窗口 reshape(-1),和同样摊平的 wf[o] 做一次 np.dot。
为什么不在最内层继续写 for c, for ki, for kj?因为那样每个标量乘加都要经过一次 Python 循环,1.51 亿次乘加会变成 1.51 亿次解释器调度,慢到无法接受。np.dot 把 144 次乘加压进一次 C 调用,已经是这个朴素框架里最划算的做法。
注意 patch 的切片:输出位置之间走 stride,窗口内部走 dilation。和 im2col 的写法正好对偶——im2col 是「固定核偏移,沿输出位置切片」,这里是「固定输出位置,沿核偏移切片」。
2.2 1.51 亿次乘加的账

<!-- 配图:三张卡片 144、1,048,576、150,994,944 用乘号与等号连接;底部左右对比 0.783 s 与 0.00796 s -->
每个位置:C×K×K = 16×9 = 144 次乘加
位置数: N×O×OH×OW = 8×32×64×64 = 1,048,576
总乘加: 144 × 1,048,576 = 150,994,944 ≈ 1.51 亿实测总时间 0.783 s,平均每个位置 746.9 ns,折合每次乘加 5.19 ns。这个数字很荒唐:一次 float32 乘加在 CPU 上应该是亚纳秒级,5.19 ns 说明时间根本不在乘加上,而在「调度这 144 次乘加」的 Python 开销上。
2.3 先用 5×5 验,再放大
在大形状上直接调参是自杀。先用案例篇的 5×5 输入跑通:
输入:
[[0, 1, 5, 1, 0],
[0, 2, 5, 2, 0],
[0, 1, 5, 1, 0],
[0, 3, 5, 3, 0],
[0, 1, 5, 1, 0]]
核:
[[−1, 0, 1],
[−1, 0, 1],
[−1, 0, 1]]
输出:
[[ 15, 0, −15],
[ 15, 0, −15],
[ 15, 0, −15]]这份输出由 assets/case_walkthrough.py 生成,用的是同一套切片逻辑。先让 3×3 的输出对上,再把 N/C/O/H/W 逐个加上去;每加一维,只验一个 1×1 的小例子。这样出错时能立刻定位是哪一维的索引写反了。
2.4 每一层到底进多少次
大案例的四层循环,进入次数是这样翻倍的:
| 层级 | 变量 | 范围 | 这一层的进入次数 |
|---|---|---|---|
| 1 | n | 0–7 | 8 |
| 2 | o | 0–31 | 8 × 32 = 256 |
| 3 | i | 0–63 | 256 × 64 = 16,384 |
| 4 | j | 0–63 | 16,384 × 64 = 1,048,576 |
| 最内层 | np.dot | 144 维 | 1,048,576 次调用 |
每往内一层,次数乘以下一维的尺寸。最内层的 np.dot 被调用了 1,048,576 次,而不是「144 次乘加被调用 1,048,576 次」——np.dot 内部还有 144 次乘加,但 Python 只看见一次函数调用。
这就是为什么把最内层从「C、K 两层循环」改成 np.dot 收益巨大:前者会让 Python 看见 1,048,576 × 144 ≈ 1.51 亿 次循环;后者只看见 1,048,576 次。1.51 亿次 Python 循环在本机上要跑几十秒,1,048,576 次只要 0.783 秒。
三、和 PyTorch 对齐:12 组配置
3.1 对齐协议
对齐不是「跑一遍看差不多」,而是固定四件事:
- 同一份
x、w、b,先转 float64,再转 float32; - 参考实现直接调
F.conv2d,不经过任何自定义封装; - 两套实现都返回同形状数组,逐元素取绝对差;
- 记录最大值,不记录平均值——最大值才能暴露边界错误。
# 第 4 章:CNN 原理/assets/run_experiments.py :: conv_ref()
def conv_ref(x, w, b, stride, padding, dilation):
return F.conv2d(torch.from_numpy(x), torch.from_numpy(w),
torch.from_numpy(b), stride=stride, padding=padding,
dilation=dilation).numpy()参考实现只有四行,作用是把「PyTorch 的答案」变成一个 numpy 数组。后面的 naive_conv2d 与 im2col_conv2d 都和它比。
12 组配置覆盖了不同的 N、C、H、W、O、K、stride、padding、dilation,包括 1×1×1×1、偶数核 4、K=7、stride 3、dilation 2。每组都跑 f64 与 f32 两遍。
3.2 12 组结果

<!-- 配图:13 行表格(表头 + 12 组),列是编号、输入输出形状、k、s、p、d、f64 最大差、f32 最大差;f64 列全部在 1e-14 量级 -->
| # | 输入 → 输出 | k | s | p | d | f64 最大差 | f32 最大差 |
|---|---|---|---|---|---|---|---|
| 1 | 2×1×5×5 → 1×3×3 | 3 | 1 | 0 | 1 | 6.66e-16 | 7.15e-07 |
| 2 | 2×3×8×8 → 4×8×8 | 3 | 1 | 1 | 1 | 7.11e-15 | 3.81e-06 |
| 3 | 1×3×9×9 → 2×5×5 | 5 | 2 | 2 | 1 | 7.11e-15 | 3.34e-06 |
| 4 | 2×4×7×7 → 3×3×3 | 3 | 1 | 0 | 2 | 2.66e-15 | 2.86e-06 |
| 5 | 1×1×1×1 → 1×1×1 | 1 | 1 | 0 | 1 | 0 | 0 |
| 6 | 3×16×14×14 → 32×7×7 | 3 | 2 | 1 | 1 | 3.91e-14 | 1.91e-05 |
| 7 | 2×32×28×28 → 64×28×28 | 3 | 1 | 1 | 1 | 9.24e-14 | 5.34e-05 |
| 8 | 1×8×10×10 → 8×10×10 | 7 | 1 | 3 | 1 | 8.53e-14 | 4.96e-05 |
| 9 | 2×2×6×6 → 2×5×5 | 2 | 1 | 0 | 1 | 8.88e-16 | 7.15e-07 |
| 10 | 1×3×12×12 → 5×4×4 | 4 | 3 | 1 | 1 | 3.55e-15 | 3.81e-06 |
| 11 | 1×1×28×28 → 1×26×26 | 3 | 1 | 0 | 1 | 1.78e-15 | 1.43e-06 |
| 12 | 2×5×11×11 → 7×5×5 | 3 | 2 | 1 | 2 | 5.33e-15 | 3.81e-06 |
f64 的 12 组全部通过,最大是第 7 组 2×32×28×28 → 64×28×28 的 9.24e-14。这个组也是 f32 最差的一组(5.34e-5),原因在 §3.5。
几个值得单独看的组:
- 第 5 组
1×1×1×1:只有一次乘加,f64 和 f32 都是 0。它证明「所有维度都等于 1」时没有隐藏的索引偏移。 - 第 8 组
K=7, p=3:奇数大核 + same padding,f64 8.53e-14。它把ki*dilation与OH*stride的边界同时压到最大。 - 第 9 组
K=2:偶数核、无 padding,输出5×5。它验证(H-K)//s+1在不能整除时不会多算一格。 - 第 10 组
K=4, s=3, p=1:偶数核 + 非整除 stride,输出4×4。它验证floor而不是ceil。 - 第 12 组
d=2, s=2:dilation 与 stride 同时出现,f64 5.33e-15。它是「起点按 dilation、步长按 stride」这条规则的压力测试。
3.3 误差落在哪个量级

<!-- 配图:对数横轴从 1e-16 到 1e-4,六条误差条:f64 9.24e-14、f32 5.34e-5、4.96e-5、1.91e-5、3.81e-6、0;红色竖线标出验收线 1e-5 -->
f64 最大: 9.24e-14
f32 最大: 5.34e-5
f32 相对最大: 9.16e-7
f32 绝对达标: 9 / 12f64 的 9.24e-14 离 1e-5 还有九个数量级,基本是「同一个公式、同样的运算顺序,只差机器舍入」。f32 的最大 5.34e-5 则已经越过验收线,但相对误差只有 9.16e-7,远小于 f32 的 eps≈1.19e-7 的 10 倍。这个对比说明:绝对误差和数值量级有关,相对误差才反映实现是否一致。
3.4 12 组覆盖了哪些边界

<!-- 配图:8 条范围条分别表示 N 1–3、C 1–32、H/W 1–28、O 1–64、K 1–7、s 1–3、p 0–3、d 1–2;底部列出 1×1×1×1、K=2/4/7、s=3、p=3、d=2 -->
| 维度 | 最小值 | 最大值 | 边界样例 |
|---|---|---|---|
| N | 1 | 3 | 单张图与 3 张图 |
| C | 1 | 32 | 单通道与 32 通道 |
| H/W | 1 | 28 | 1×1 与 28×28 |
| O | 1 | 64 | 单核与 64 核 |
| K | 1 | 7 | 1×1、2、4(偶数)、7 |
| stride | 1 | 3 | 非整除的 3 |
| padding | 0 | 3 | same 与 valid |
| dilation | 1 | 2 | 空洞卷积 |
覆盖不是越多越好,而是要覆盖「会让索引公式出错」的组合:偶数核、非整除 stride、dilation + stride 同时出现、1×1 输入、K=1。12 组里这五类都有。
3.5 踩坑:f64 全过,f32 三组超线

<!-- 配图:三栏「现象 / 原因 / 修复」,分别列出 3/12 组超标、f32 尾数与累加顺序、改用 f64 或相对误差 -->
- 现象:同一份代码、同一组输入,f64 12/12 通过,最大 9.24e-14;换成 f32 后最大 5.34e-5,第 6、7、8 组超过 1e-5。第 7 组是最大的形状
2×32×28×28 → 64×28×28。 - 原因:float32 只有 24 位尾数,机器精度
eps≈1.19e-7。每个输出是 144 项乘积之和,累加顺序与 PyTorch 的 FMA 路径不同,误差随项数和数值量级增长。第 7 组的输出量级到 ~12,绝对误差被放大到 5e-5 是正常的;相对误差仍是 9.16e-7。 - 修复:验收用 float64,不用 f32;如果必须在 f32 下比较,用
atol=1e-5, rtol=1e-4或直接看相对误差,不要用一个固定的绝对阈值判所有形状。
这条踩坑的结论要记牢:「逐元素误差 < 1e-5」是一个带 dtype 和量级的判据,不是一个放之四海皆准的常数。本篇的正式验收因此写成「float64 下 12/12 < 1e-5」;f32 只作为压力测试,并额外报告相对误差。
3.6 对不上时,按这个顺序查
12 组都通过不代表你改完代码还能通过。真正写自己的卷积时,对不上是常态。排查顺序不能反:先形状,后数值。形状不对时,比较数值没有意义——两个数组的元素根本不在同一位置。
| 症状 | 先查 | 典型原因 | 怎么确认 |
|---|---|---|---|
| 形状不一致 | 输出尺寸公式 | stride/padding/dilation 传错位置 | 手算 OH/OW,和 im2col() 返回的两个数比 |
| 形状一致,整张图差一个倍数 | 权重布局 | w 是 (O,C,K,K),却被当成 (C,O,K,K) | 打印 w.shape 与 w.reshape(O,-1)[0,:6] |
| 形状一致,只有边缘错 | padding | 补零宽度或 same 的 p 算错 | 先把 p=0 的 valid 配置跑通,再加 padding |
| 形状一致,误差随数值变大 | dtype | f32 累加舍入 | 同一组输入换 f64,看最大差是否降到 1e-13 |
| 形状一致,个别位置差很大 | 切片起点 | dilation 只改起点,stride 只改步长 | 打印一组 (ki,kj) 对应的输入下标 |
前三条是「索引错」,后两条是「数值错」。索引错会在形状或边缘立刻暴露;数值错通常整张图都「差不多对」,只是尾数或个别点不对。把这两类分开,能少走很多弯路。
一组配置对不上时,先跑这个:
# 第 4 章:CNN 原理/assets/run_experiments.py :: 单组排查(临时脚本)
import numpy as np
from run_experiments import conv_ref, im2col_conv2d
def check_one(x, w, b, s=1, p=0, d=1):
ref = conv_ref(x, w, b, s, p, d)
mine = im2col_conv2d(x, w, b, s, p, d)
print("形状", ref.shape, mine.shape)
if ref.shape != mine.shape:
return
diff = np.abs(ref - mine)
k = np.unravel_index(diff.argmax(), diff.shape)
print("最大绝对差", diff.max(), "位置", k)
print("该点 ref/mine", ref[k], mine[k])
print("最大相对差", (diff / (np.abs(ref) + 1e-12)).max())它一次打印四样东西:形状、最大绝对差、最差位置的坐标、该点两边的值。读法也有顺序:
- 形状不同,直接回去查
out_size()的三个参数。此时不要看后面的数。 - 形状相同但最差位置在角落,八成是 padding;在中间且只有零星几个点,八成是切片起点。
- 最差位置的值一个正常、一个数量级离谱,是权重布局或索引错位,不是精度问题。
- 最大绝对差在 1e-13 附近、最大相对差也在 1e-13 附近,就是通过;绝对差大但相对差小,是 f32 的尺度效应,不是实现错。
run_experiments.py 的 main() 在 if __name__ == "__main__" 里,所以 import 它不会触发那 30 秒的完整实验。排查脚本可以反复跑,不会覆盖 metrics.json。
四、把 conv2d 变成可训练层
数值对齐只证明前向。要训练,还得让梯度也走同一条路。PyTorch 的自定义 autograd.Function 要求你手写 backward,这正好强迫你把卷积的反向公式推一遍。
4.1 前向四步,反向四公式

<!-- 配图:上排前向 x→unfold→cols→matmul→+bias→y;下排红色虚线反向 grad_y→Wᵀ·go→grad_cols→F.fold→grad_x;底部列出 grad_w 与 grad_b 的公式 -->
前向:
cols = F.unfold(x) # (N, CKK, L)
z = W.view(O, CKK) @ cols # (N, O, L)
y = z.view(N, O, OH, OW) + b # 广播到每个位置反向,已知 go = grad_y:
grad_cols = Wᵀ @ go # (N, CKK, L)
grad_x = F.fold(grad_cols) # 重叠位置相加
grad_w = Σ_N (go @ colsᵀ) # (O, CKK)
grad_b = Σ_{N,H,W} go # (O,)F.fold 是 F.unfold 的逆操作。前向 unfold 会把重叠的输入复制到多列,反向 fold 就把多列的梯度加回同一个输入位置。这不是近似,是卷积反向的定义:一个输入像素被 9 个窗口用过,它的梯度就是这 9 条路径之和。
4.2 代码
# 第 4 章:CNN 原理/assets/run_experiments.py :: _Conv2dFn
class _Conv2dFn(torch.autograd.Function):
@staticmethod
def forward(ctx, x, w, b, stride, padding, dilation):
N, C, H, W = x.shape
O, _, K, _ = w.shape
OH, OW = out_size(H, K, stride, padding, dilation), out_size(W, K, stride, padding, dilation)
cols = F.unfold(x, (K, K), dilation=dilation, padding=padding, stride=stride)
ctx.save_for_backward(x, w, b, cols)
ctx.stride, ctx.padding, ctx.dilation = stride, padding, dilation
out = torch.matmul(w.view(O, -1), cols).view(N, O, OH, OW)
return out + b.view(1, O, 1, 1)
@staticmethod
def backward(ctx, grad_out):
x, w, b, cols = ctx.saved_tensors
N, C, H, W = x.shape
O, _, K, _ = w.shape
go = grad_out.contiguous().view(N, O, -1)
grad_cols = torch.matmul(w.view(O, -1).t(), go)
grad_x = F.fold(grad_cols, output_size=(H, W), kernel_size=(K, K),
dilation=ctx.dilation, padding=ctx.padding, stride=ctx.stride)
grad_w = torch.matmul(go, cols.transpose(1, 2)).sum(0).view(O, C, K, K)
grad_b = grad_out.sum(dim=(0, 2, 3))
return grad_x, grad_w, grad_b, None, None, Noneforward 里 torch.matmul(w.view(O, -1), cols) 是 (O,CKK)×(N,CKK,L),PyTorch 会广播成 (N,O,L)。backward 返回 6 个值,对应 forward 的 6 个输入;后三个 None 是 stride/padding/dilation,它们没有梯度。
ctx.save_for_backward 只存 x、w、b、cols,不存 out。grad_w 的计算需要 cols,grad_b 只需要 grad_out,grad_x 需要 w 与 grad_out。存什么、不存什么是显存与计算的取舍:存 cols 省一次 unfold,代价是 18 MB 级别的列矩阵一直挂在计算图上。
4.3 梯度对齐
把自定义层和 nn.Conv2d 设成同一份权重,分别对 x、w、b 求梯度:

<!-- 配图:四张卡片:前向 4.44e-16、grad_x 0、grad_w 1.24e-14、grad_b 0;底部绿色 pass 徽章,判据 <1e-10 -->
| 项目 | 最大绝对差 |
|---|---|
| 前向输出 | 4.44e-16 |
| grad_x | 0 |
| grad_w | 1.24e-14 |
| grad_b | 0 |
四项都远低于 1e-10 的判据。grad_x 与 grad_b 直接是 0,说明 F.fold 的重叠加和与 sum 没有引入额外误差;grad_w 的 1.24e-14 来自两次矩阵乘的舍入,量级和 f64 的机器精度一致。
这里没有「大概对」:pass 的条件是 max(fwd, grad_x, grad_w, grad_b) < 1e-10,写在 exp_grad() 里,不通过就抛错。
4.4 接进 nn.Module
# 第 4 章:CNN 原理/assets/run_experiments.py :: Im2ColConv2d
class Im2ColConv2d(nn.Module):
"""前向 unfold + 矩阵乘,反向 fold + 矩阵乘。"""
def __init__(self, c_in, c_out, k=3, stride=1, padding=0, dilation=1, bias=True):
super().__init__()
self.weight = nn.Parameter(torch.empty(c_out, c_in, k, k))
self.bias = nn.Parameter(torch.empty(c_out)) if bias else None
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
bound = 1 / math.sqrt(c_in * k * k)
nn.init.uniform_(self.bias, -bound, bound)
self.stride, self.padding, self.dilation = stride, padding, dilation
def forward(self, x):
return _Conv2dFn.apply(x, self.weight, self.bias,
self.stride, self.padding, self.dilation)它和 nn.Conv2d 的接口一致:c_in、c_out、k、stride、padding、dilation、bias。把它塞进第 3 章的小 CNN:
# 第 4 章:CNN 原理/assets/run_experiments.py :: SmallCNN
class SmallCNN(nn.Module):
def __init__(self, conv_cls=Im2ColConv2d):
super().__init__()
self.c1 = conv_cls(1, 16, 3, padding=1)
self.c2 = conv_cls(16, 32, 3, padding=1)
self.fc = nn.Linear(32 * 7 * 7, 10)
def forward(self, x):
x = F.max_pool2d(F.relu(self.c1(x)), 2)
x = F.max_pool2d(F.relu(self.c2(x)), 2)
x = x.flatten(1)
return self.fc(x)训练结果:

<!-- 配图:三张卡片 98.52%、20,490、16.80 s;下方三行 epoch 表列出 train_acc 与 val_acc;底部说明模块前向最大差 4.77e-6 -->
| 项目 | 数值 |
|---|---|
| 测试集准确率(3 轮) | 98.52% |
| 参数量 | 20,490(MLP 对照 50,890) |
| 训练耗时 | 16.80 s(MLP 0.32 s) |
| 与 nn.Conv2d 模块前向最大差 | 4.77e-6 |
第 3 章的 MLP 在同一份数据上测试 96.04%;这个用自写卷积层的小 CNN 到 98.52%,高了 2.48 个百分点。参数量只有 20,490,是 MLP 50,890 的 40%。训练能收敛、准确率正常,说明 backward 不只是「数值上对」,而是真的在优化。
训练耗时 16.80 s 比 MLP 的 0.32 s 慢 52 倍,这并不奇怪:小 CNN 每张图的乘加是 1,031,744,MLP 是 50,816,相差 20.3 倍;再加上 unfold/fold 的展开与折叠,慢 52 倍在同一个量级。
4.5 设计取舍:自己写 autograd,还是直接用 F.conv2d
| 方案 | 收益 | 代价 |
|---|---|---|
A:自定义 _Conv2dFn | 前向/反向每一步可见;可以插调试、改 unfold 策略;理解卷积反向 | 要自己推公式、维护 backward;f32 前向差 4.77e-6;比 torch 慢 14.7× |
B:直接调 F.conv2d | 最快;梯度由框架维护;数值稳定 | 黑盒,看不到 im2col 与 fold;想改内部路径只能等上游 |
这一篇选 A,因为目标是把机制看清楚;生产代码选 B,因为 torch 的卷积内核比 Python 侧的 im2col 快一个数量级。两者不是替代关系:A 是教学与排查工具,B 是部署工具。你可以用 A 生成一份「参考答案」,再用 B 跑大训练。
4.6 反向公式的一个标量例子
§4.2 的三个矩阵公式看得懂,却不容易「感觉到」。把它们缩小到一个乘加:
前向:y = x·w + b = 2×3 + 1 = 7
上游:grad_y = 1链式法则只做一件事:每个输入的梯度 = 上游梯度 × 本地偏导。
| 变量 | 前向值 | 本地偏导 | 梯度 |
|---|---|---|---|
| x | 2 | ∂y/∂x = w = 3 | 1×3 = 3 |
| w | 3 | ∂y/∂w = x = 2 | 1×2 = 2 |
| b | 1 | ∂y/∂b = 1 | 1×1 = 1 |
用代码把这三个数算出来:
# 第 4 章:CNN 原理/assets/run_experiments.py :: exp_backward_scalar()
x, w, b, go = 2.0, 3.0, 1.0, 1.0
y = x * w + b
grad_x, grad_w, grad_b = go * w, go * x, go
print(y, grad_x, grad_w, grad_b) # 7.0 3.0 2.0 1.0这三个数就是 §4.2 三个矩阵公式在「只有一个位置、一个通道」时的样子:
grad_b = Σ go:标量里只有一项,所以就是 1;矩阵里对N、OH、OW求和,因为同一个b[o]被每个输出位置用过。grad_w = Σ (go × cols):标量里是1×2;矩阵里对N和所有输出位置求和,因为同一个w[o,c,ki,kj]在每个位置都被用过一次。grad_x = go × w:标量里是1×3;矩阵里一个输入像素会被最多K×K个窗口用到,所以F.fold必须把这些路径的梯度加回去。
一句话记法:反向就是把上游梯度按「前向用过的次数」分摊回去。w 和 b 在每个输出位置都用过,所以求和;x 被多个重叠窗口用过,所以 fold 相加。没有额外的魔法。
五、速度:慢在哪
5.1 三档尺寸

<!-- 配图:9 条横向对数柱,按小/中/大分组,每组 naive(红)、im2col(蓝)、torch(绿);底部给出大案例 0.783 s → 0.00796 s → 0.000543 s -->
| 档位 | 形状 | naive | im2col | torch | naive→im2col |
|---|---|---|---|---|---|
| 小 | 1×1×16×16 → 8×16×16 | 1.279 ms | 0.03858 ms | 0.01517 ms | 33.1× |
| 中 | 4×8×32×32 → 16×32×32 | 45.27 ms | 0.3102 ms | 0.06025 ms | 146.0× |
| 大 | 8×16×64×64 → 32×64×64 | 783.2 ms | 7.958 ms | 0.5426 ms | 98.4× |
三档的加速比不是单调的:小案例只有 33.1×,因为矩阵太小,BLAS 的多线程还没热身;中案例 146.0×,形状刚好吃满;大案例 98.4×,列矩阵 18 MB 开始受内存带宽影响。这说明「im2col 一定快 100 倍」是错的,快多少取决于形状。
GFLOPS 也一样:
小:0.029 / 0.955 / 2.431
中:0.208 / 30.426 / 156.634
大:0.386 / 37.948 / 556.578朴素实现的 GFLOPS 始终在 0.4 以下,im2col 到 38,torch 到 557。torch 在大案例上是 im2col 的 14.7 倍,因为它不物化列矩阵,直接把卷积映射成内核里的隐式 GEMM。
5.2 K 扫描:乘加涨 25 倍,朴素时间只涨 1.41 倍

<!-- 配图:折线图 K=1/3/5,三条线分别是乘加量 1/9/25、im2col 时间 1/10.93/39.98、朴素时间 1/1.31/1.41;底部说明朴素时间被固定循环开销吃掉 -->
| K | 乘加量 | naive 时间 | im2col 时间 | 乘加比 | naive 比 | im2col 比 |
|---|---|---|---|---|---|---|
| 1 | 2,097,152 | 0.1375 s | 0.0001375 s | 1 | 1 | 1 |
| 3 | 18,874,368 | 0.1796 s | 0.001504 s | 9 | 1.306 | 10.93 |
| 5 | 52,428,800 | 0.1941 s | 0.005499 s | 25 | 1.412 | 39.98 |
这张表是「慢在解释器,不在数学」的直接证据:
- 乘加量从 1 涨到 25 倍,朴素时间只涨 1.41 倍。多出来的 24 倍乘加被 262,144 次固定的 Python 循环开销吞掉了,乘加本身不是瓶颈。
- im2col 时间涨 39.98 倍,略高于乘加量的 25 倍。多出来的部分来自列矩阵变大后的内存访问:K=5 时列高 200,
cols矩阵比 K=1 大 25 倍,写列与读列都更贵。 - 朴素时间在 K=1 时就已经是 0.1375 s,而 im2col 只要 0.0001375 s,差 1000 倍。这不是 K 的问题,是「每个位置一次 Python 调度」的问题。
5.3 慢在哪:746.9 ns 对 7.59 ns

<!-- 配图:三条横向对数条,朴素循环 746.9 ns/位置、im2col 7.59 ns/位置、torch 0.517 ns/位置;右侧注明各自的原因;底部绿色验收条 -->
朴素: 0.783 s / 1,048,576 = 746.9 ns/位置
im2col: 0.00796 s / 1,048,576 = 7.59 ns/位置
torch: 0.000543 s / 1,048,576 = 0.517 ns/位置三个数字对应三种调度方式:
- 朴素循环:每个输出位置都要经过一次 Python 的
for、一次切片、一次reshape、一次np.dot调用。746.9 ns 里,真正做 144 次乘加的时间不到 1 ns,其余都是解释器与 NumPy 的调用开销。 - im2col + BLAS:1,048,576 个位置先被展开成矩阵,然后一次
matmul交给多线程 BLAS。7.59 ns/位置 已经接近「内存搬运 + 乘加」的成本,Python 只调度一次矩阵乘。 - torch 内核:0.517 ns/位置,连列矩阵都不物化。oneDNN/cuDNN 直接在输入上做隐式 GEMM,寄存器里完成乘加,内存只读一遍。
所以「手写版本比 PyTorch 慢 99 倍」的准确说法是:慢在 1,048,576 次 Python 位置循环,不在 1.51 亿次乘加。把循环换成矩阵乘,就能从 0.783 s 降到 0.00796 s;再往下要降到 0.000543 s,得换 C++/CUDA 内核,而不是继续优化 Python。
5.4 设计取舍:可读性还是速度
| 方案 | 收益 | 代价 |
|---|---|---|
| A:朴素循环 | 代码和公式一一对应;无额外显存;改参数只改切片 | 大案例 0.783 s;GFLOPS 0.386;不能用于真实训练 |
| B:im2col + BLAS | 0.00796 s,98.4×;GFLOPS 37.9;代码仍然是 Python | 18 MB 列矩阵;9 倍显存;需要矩阵库 |
这一篇两个都保留:A 放在 naive_conv2d(),作为对齐与教学的基线;B 放在 im2col_conv2d(),作为可用的快速路径。真实框架再往前一步,用 implicit GEMM 把 B 的列矩阵也省掉,但那是 §5.3 的第三层,不是这一篇的范围。
5.5 三个版本放一起
| 版本 | 位置 | 速度(大案例) | 显存 | 用途 |
|---|---|---|---|---|
| 朴素四重循环 | naive_conv2d() | 0.783 s | 输入 + 输出 | 正确性基线、教学 |
| im2col + BLAS | im2col_conv2d() | 0.00796 s | 输入 + 18 MB cols | 快速验证、中小模型 |
| torch 内核 | F.conv2d | 0.000543 s | 输入 + 输出 | 生产训练 |
5.6 如果把 im2col 用在真实网络
大案例的列矩阵 18 MB 看起来还能接受,因为输入只有 8×16×64×64。换成真实网络的中间层:单张 64×224×224、K=3、same、stride 1。
- 输入一张:
64×224×224×4= 12.25 MB - 列矩阵一张:
64×9×224×224×4= 110.25 MB,仍然是输入的 9 倍 - batch 32:光这一层的
cols就是 3528 MB,约 3.4 GB
这还只是前向。§4.2 为了算 grad_w 把 cols 存进了 ctx.save_for_backward,意味着这 3.4 GB 要一直挂到反向;再加上激活值、梯度和优化器状态,一张 24 GB 的卡跑一层就满了。所以 im2col 不能直接用在真实训练里。
真实框架走的是 implicit GEMM:不物化列矩阵,把输入切成能装进缓存的小块,在寄存器里一边取数一边乘加。§5.1 里 torch 比 im2col 快 14.7 倍,差的就是这 3.4 GB 的搬运和写入。
结论不是「im2col 没用」,而是它的位置是诊断工具:当形状对不上、想知道某个输出位置到底取了哪些输入时,展开成矩阵最直观;当要跑大模型时,把这套索引交给框架的内核。会写 im2col 的人,看 unfold、fold、groups、dilation 的参数错误会快很多,但不会真的在生产代码里 np.empty 一个 3.4 GB 的列矩阵。
5.7 测速口径
上面的毫秒数都来自同一次运行,口径固定:
- 计时函数
timeit(fn, repeats, warmup=1):先空跑 1 次热身,再重复repeats次,取中位数,避免单次抖动。 - 朴素大案例只跑 1 次(0.783 s 一次就够);小、中各 3 次;im2col 与 torch 各 5 次。
- 固定 4 线程、float32、同一进程;输入数据在计时前生成,不计入。
- 只计函数调用的墙钟时间,不含首次 import、JIT 或显存分配。
所以别把 0.783 s 当成绝对值。换机器、换线程数,毫秒会变;能复现的是比例:三档的 33.1× / 146.0× / 98.4×,以及 K 扫描的 1.41× 与 39.98×。报告速度时,比例比绝对值更可信。
六、验收:把三条标准跑一遍
把这一篇的交付物逐条对照:
| 验收项 | 标准 | 实测 | 结果 |
|---|---|---|---|
| f64 数值对齐 | 12 组逐元素最大差 < 1e-5 | 9.24e-14 | 通过 |
| 边界形状 | 1×1、K=1/2/4/7、s=3、p=3、d=2 | 12 组全部通过 | 通过 |
| f32 压力测试 | 记录而不是隐藏 | 5.34e-5,9/12 绝对达标 | 已解释 |
| 自定义层前向 | < 1e-10 | 4.44e-16 | 通过 |
| 自定义层梯度 | grad_x/grad_w/grad_b < 1e-10 | 最大 1.24e-14 | 通过 |
| 能训练 | 3 轮 MNIST | 98.52% | 通过 |
| 能解释慢 | 说出瓶颈 | Python 位置循环 746.9 ns | 通过 |
用一段断言把前向对齐钉死:
# 第 4 章:CNN 原理/assets/run_experiments.py :: exp_align() 的核心断言
ref64 = conv_ref(x, w, b, s, p, d)
mine64 = im2col_conv2d(x, w, b, s, p, d)
assert np.max(np.abs(ref64 - mine64)) < 1e-12
ref32 = conv_ref(x.astype(np.float32), w.astype(np.float32),
b.astype(np.float32), s, p, d)
mine32 = im2col_conv2d(x.astype(np.float32), w.astype(np.float32),
b.astype(np.float32), s, p, d)
assert np.max(np.abs(ref32 - mine32)) < 6e-5第一条断言是正式验收(f64 <1e-12,比 1e-5 更严);第二条是 f32 的宽松上界,只用于回归,不用于「证明实现正确」。
这三条标准防的是三种不同的错:
- 数值对齐防「索引写错」:形状对了、值不对,通常是切片起点或权重布局。
- 梯度对齐防「反向写错」:前向对、梯度错,通常是 fold 忘了相加,或
grad_w少了一个转置。 - 速度基准防「解释错了」:只看到慢,不知道慢在 Python 调度还是乘加,就没法优化。
任何一条单独通过都不够:只做前向的实现可以数值全对但无法训练;梯度对的实现可以慢到不可用;快的实现可以悄悄算错。三条都过,才叫闭环。
到这一步,卷积的前向、反向、训练、速度都闭环了。但还有一个这一章解决不了的问题:
网络能堆多深?
理论篇结尾提过:把 3×3 卷积一层层堆起来,感受野变大、参数不变,看起来堆得越深越好。但真实训练里,层数一多,训练误差反而上升——不是过拟合,是优化本身变难。ResNet 用 152 层、残差连接解决了这个问题,但那属于第 5 章。这一篇的 conv2d 只是砖块,砖块本身不会告诉你「为什么 152 层能训、56 层反而更差」。
自检
合上笔记,试着回答:
- im2col 的列高和列数分别由什么决定?
C_in=16、K=3、H_out=W_out=64、N=8时,列矩阵是多大? - 为什么 im2col 的切片步长是
stride,而 dilation 只改起点?两者写反会得到什么形状? - 朴素四重循环里,最内层为什么用
np.dot而不是继续写C、K两层? - 12 组对齐里,哪三组 f32 超过 1e-5?为什么 f64 全过?
F.fold在反向里起什么作用?如果去掉它,grad_x会变成什么形状?- 自定义层的
backward返回几个值?后三个为什么是None? - 大案例的 0.783 s 主要花在哪里?为什么 K 从 1 到 5,朴素时间只涨 1.41 倍?
- 这一篇的验收为什么要分 f64 与 f32 两套口径?
- 一组配置形状相同、最差位置在角落,先查哪一项?最差位置在中间且只有零星几个点,又先查什么?
y = x·w + b的标量反向里,grad_x为什么是 3 而不是 1?- 单张
64×224×224的列矩阵为什么是 110.25 MB?batch 32 时这个数会变成多少?