Skip to content

实战演练:手写 conv2d 并与 PyTorch 对齐 ​

案例篇把 5×5 的 81 次乘法算到了每一位。真实网络里,一层 3×3 卷积就是 22 万次乘加,一个 MNIST 小 CNN 每张图 103 万次。这一篇把「算一次」变成「算一亿次」:先写 im2col,再写朴素四重循环,再和 F.conv2d 逐元素对齐,最后把卷积包成可训练层,量出慢在哪。

所有代码在 assets/run_experiments.py,所有数字由它写入 assets/metrics.json。从头跑一遍:

bash
cd "第 4 章:CNN 原理/assets"
/private/tmp/dnn-venv/bin/python run_experiments.py

本机 29.9 秒结束,输出 12 组对齐、4 项梯度、3 档速度基准。

![实战篇的五个阶段](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_roadmap.svg)

<!-- 配图:五张卡片从左到右排列:im2col、朴素循环、12 组对齐、自定义层、速度基准;每张卡片底部给出产物或关键数字,最后一张 98.4× -->

零、先把「对」写死 ​

「和 PyTorch 对齐」不是一句感觉,而是三条可执行的判据。先写判据,再写代码,后面所有优化才有裁判。

![三条验收线](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_accept_target.svg)

<!-- 配图:三张卡片分别给出数值 1e-5、梯度 1e-10、速度 98.4×;底部环境条列出 Python 3.9.6、torch 2.8.0、numpy 2.0.2、4 线程 -->

项目判据实测
数值(float64)12 组输出逐元素最大差 < 1e-59.24e-14
数值(float32,附加)同上,仅作压力测试5.34e-5,9/12 达标
梯度前向、grad_x、grad_w、grad_b 最大差 < 1e-101.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 的小块。把这些小块按「输出位置」排成列,就把卷积变成了矩阵乘。

![im2col 把一个窗口摊成一列](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_im2col.svg)

<!-- 配图:左边 4×4 输入高亮左上 3×3 窗口;右边 9×4 矩阵,第一列就是该窗口的 9 个数,其余三列对应另外三个窗口 -->

4×4 输入、3×3 核、stride 1,输出是 2×2,一共 4 个位置。把每个位置的 9 个输入按 (通道, 核行, 核列) 的顺序排好,就得到 9 行 4 列:

text
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 不需要四层循环。它只遍历「通道、核行、核列」三层,每一层把一个固定偏移的采样面切片、拉平、塞进一行:

python
# 第 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:

![矩阵乘的维度](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_im2col_shape.svg)

<!-- 配图:W 为 4×9、cols 为 9×6、out 为 4×6;乘号与等号连接三个矩阵,黄色格标出 out(0,0) 由 W 第 0 行与 cols 第 0 列点积得到 -->

text
W:      O × CKK     4 × 9
cols:   CKK × L     9 × 6
out:    O × L       4 × 6

CKK 必须同时是 W 的列数和 cols 的行数。真实大案例是 32×144 · 144×32,768 = 32×32,768:一次矩阵乘,把 1,048,576 个输出位置的 144 次乘加全部交给 BLAS。

代码只有一行核心:

python
# 第 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 卷积的四个阶段](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_pipeline.svg)

<!-- 配图:输入 → im2col → matmul → reshape(+bias) → 输出五个节点,im2col 与 matmul 用蓝框;底部三条说明顺序不能换 -->

text
输入 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 倍:

![im2col 的 9 倍显存代价](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_memory.svg)

<!-- 配图:左边一张 2 MB 输入方块,右边 3×3 共 9 张同样大小的方块代表 18 MB 列矩阵;底部三张卡片给出 2.0 MB、18.0 MB、9.0× -->

大案例的账:

text
输入:  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×;矩阵乘交给 BLAS18 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 完全一致。

python
# 第 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 四层循环 ​

先写最笨的版本。它不追求快,只追求「每一笔都能和公式对上」:

![朴素实现的四层循环](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_naive_loop.svg)

<!-- 配图:四层嵌套框 N→O→行→列,最内层黄色区域给出 patch 切片与 np.dot 一行;底部总计 1,048,576 次最内层、150,994,944 次乘加 -->

python
# 第 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 亿次乘加的账 ​

![朴素实现的成本](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_naive_cost.svg)

<!-- 配图:三张卡片 144、1,048,576、150,994,944 用乘号与等号连接;底部左右对比 0.783 s 与 0.00796 s -->

text
每个位置: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 输入跑通:

text
输入:
[[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 每一层到底进多少次 ​

大案例的四层循环,进入次数是这样翻倍的:

层级变量范围这一层的进入次数
1n0–78
2o0–318 × 32 = 256
3i0–63256 × 64 = 16,384
4j0–6316,384 × 64 = 1,048,576
最内层np.dot144 维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 对齐协议 ​

对齐不是「跑一遍看差不多」,而是固定四件事:

  1. 同一份 x、w、b,先转 float64,再转 float32;
  2. 参考实现直接调 F.conv2d,不经过任何自定义封装;
  3. 两套实现都返回同形状数组,逐元素取绝对差;
  4. 记录最大值,不记录平均值——最大值才能暴露边界错误。
python
# 第 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 组结果 ​

![12 组对齐结果](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_align_table.svg)

<!-- 配图:13 行表格(表头 + 12 组),列是编号、输入输出形状、k、s、p、d、f64 最大差、f32 最大差;f64 列全部在 1e-14 量级 -->

#输入 → 输出kspdf64 最大差f32 最大差
12×1×5×5 → 1×3×331016.66e-167.15e-07
22×3×8×8 → 4×8×831117.11e-153.81e-06
31×3×9×9 → 2×5×552217.11e-153.34e-06
42×4×7×7 → 3×3×331022.66e-152.86e-06
51×1×1×1 → 1×1×1110100
63×16×14×14 → 32×7×732113.91e-141.91e-05
72×32×28×28 → 64×28×2831119.24e-145.34e-05
81×8×10×10 → 8×10×1071318.53e-144.96e-05
92×2×6×6 → 2×5×521018.88e-167.15e-07
101×3×12×12 → 5×4×443113.55e-153.81e-06
111×1×28×28 → 1×26×2631011.78e-151.43e-06
122×5×11×11 → 7×5×532125.33e-153.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 误差落在哪个量级 ​

![误差尺度](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_align_error.svg)

<!-- 配图:对数横轴从 1e-16 到 1e-4,六条误差条:f64 9.24e-14、f32 5.34e-5、4.96e-5、1.91e-5、3.81e-6、0;红色竖线标出验收线 1e-5 -->

text
f64 最大:        9.24e-14
f32 最大:        5.34e-5
f32 相对最大:    9.16e-7
f32 绝对达标:    9 / 12

f64 的 9.24e-14 离 1e-5 还有九个数量级,基本是「同一个公式、同样的运算顺序,只差机器舍入」。f32 的最大 5.34e-5 则已经越过验收线,但相对误差只有 9.16e-7,远小于 f32 的 eps≈1.19e-7 的 10 倍。这个对比说明:绝对误差和数值量级有关,相对误差才反映实现是否一致。

3.4 12 组覆盖了哪些边界 ​

![12 组配置的覆盖范围](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_cover.svg)

<!-- 配图: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 -->

维度最小值最大值边界样例
N13单张图与 3 张图
C132单通道与 32 通道
H/W1281×1 与 28×28
O164单核与 64 核
K171×1、2、4(偶数)、7
stride13非整除的 3
padding03same 与 valid
dilation12空洞卷积

覆盖不是越多越好,而是要覆盖「会让索引公式出错」的组合:偶数核、非整除 stride、dilation + stride 同时出现、1×1 输入、K=1。12 组里这五类都有。

3.5 踩坑:f64 全过,f32 三组超线 ​

![f32 踩坑](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_f32_trap.svg)

<!-- 配图:三栏「现象 / 原因 / 修复」,分别列出 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
形状一致,误差随数值变大dtypef32 累加舍入同一组输入换 f64,看最大差是否降到 1e-13
形状一致,个别位置差很大切片起点dilation 只改起点,stride 只改步长打印一组 (ki,kj) 对应的输入下标

前三条是「索引错」,后两条是「数值错」。索引错会在形状或边缘立刻暴露;数值错通常整张图都「差不多对」,只是尾数或个别点不对。把这两类分开,能少走很多弯路。

一组配置对不上时,先跑这个:

python
# 第 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 前向四步,反向四公式 ​

![自定义层的前向与反向](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_autograd.svg)

<!-- 配图:上排前向 x→unfold→cols→matmul→+bias→y;下排红色虚线反向 grad_y→Wᵀ·go→grad_cols→F.fold→grad_x;底部列出 grad_w 与 grad_b 的公式 -->

前向:

text
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:

text
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 代码 ​

python
# 第 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, None

forward 里 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 求梯度:

![梯度对齐](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_grad.svg)

<!-- 配图:四张卡片:前向 4.44e-16、grad_x 0、grad_w 1.24e-14、grad_b 0;底部绿色 pass 徽章,判据 <1e-10 -->

项目最大绝对差
前向输出4.44e-16
grad_x0
grad_w1.24e-14
grad_b0

四项都远低于 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 ​

python
# 第 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:

python
# 第 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)

训练结果:

![自定义层训练结果](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_train.svg)

<!-- 配图:三张卡片 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 的三个矩阵公式看得懂,却不容易「感觉到」。把它们缩小到一个乘加:

text
前向:y = x·w + b = 2×3 + 1 = 7
上游:grad_y = 1

链式法则只做一件事:每个输入的梯度 = 上游梯度 × 本地偏导。

变量前向值本地偏导梯度
x2∂y/∂x = w = 31×3 = 3
w3∂y/∂w = x = 21×2 = 2
b1∂y/∂b = 11×1 = 1

用代码把这三个数算出来:

python
# 第 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 三档尺寸 ​

![三档尺寸的速度对比](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_bench.svg)

<!-- 配图:9 条横向对数柱,按小/中/大分组,每组 naive(红)、im2col(蓝)、torch(绿);底部给出大案例 0.783 s → 0.00796 s → 0.000543 s -->

档位形状naiveim2coltorchnaive→im2col
小1×1×16×16 → 8×16×161.279 ms0.03858 ms0.01517 ms33.1×
中4×8×32×32 → 16×32×3245.27 ms0.3102 ms0.06025 ms146.0×
大8×16×64×64 → 32×64×64783.2 ms7.958 ms0.5426 ms98.4×

三档的加速比不是单调的:小案例只有 33.1×,因为矩阵太小,BLAS 的多线程还没热身;中案例 146.0×,形状刚好吃满;大案例 98.4×,列矩阵 18 MB 开始受内存带宽影响。这说明「im2col 一定快 100 倍」是错的,快多少取决于形状。

GFLOPS 也一样:

text
小: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 扫描](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_k_scan.svg)

<!-- 配图:折线图 K=1/3/5,三条线分别是乘加量 1/9/25、im2col 时间 1/10.93/39.98、朴素时间 1/1.31/1.41;底部说明朴素时间被固定循环开销吃掉 -->

K乘加量naive 时间im2col 时间乘加比naive 比im2col 比
12,097,1520.1375 s0.0001375 s111
318,874,3680.1796 s0.001504 s91.30610.93
552,428,8000.1941 s0.005499 s251.41239.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 ​

![慢在哪](/ai-four-kings/第 4 章:CNN 原理/assets/ch4_pro_slow.svg)

<!-- 配图:三条横向对数条,朴素循环 746.9 ns/位置、im2col 7.59 ns/位置、torch 0.517 ns/位置;右侧注明各自的原因;底部绿色验收条 -->

text
朴素:        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/位置

三个数字对应三种调度方式:

  1. 朴素循环:每个输出位置都要经过一次 Python 的 for、一次切片、一次 reshape、一次 np.dot 调用。746.9 ns 里,真正做 144 次乘加的时间不到 1 ns,其余都是解释器与 NumPy 的调用开销。
  2. im2col + BLAS:1,048,576 个位置先被展开成矩阵,然后一次 matmul 交给多线程 BLAS。7.59 ns/位置 已经接近「内存搬运 + 乘加」的成本,Python 只调度一次矩阵乘。
  3. 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 + BLAS0.00796 s,98.4×;GFLOPS 37.9;代码仍然是 Python18 MB 列矩阵;9 倍显存;需要矩阵库

这一篇两个都保留:A 放在 naive_conv2d(),作为对齐与教学的基线;B 放在 im2col_conv2d(),作为可用的快速路径。真实框架再往前一步,用 implicit GEMM 把 B 的列矩阵也省掉,但那是 §5.3 的第三层,不是这一篇的范围。

5.5 三个版本放一起 ​

版本位置速度(大案例)显存用途
朴素四重循环naive_conv2d()0.783 s输入 + 输出正确性基线、教学
im2col + BLASim2col_conv2d()0.00796 s输入 + 18 MB cols快速验证、中小模型
torch 内核F.conv2d0.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-59.24e-14通过
边界形状1×1、K=1/2/4/7、s=3、p=3、d=212 组全部通过通过
f32 压力测试记录而不是隐藏5.34e-5,9/12 绝对达标已解释
自定义层前向< 1e-104.44e-16通过
自定义层梯度grad_x/grad_w/grad_b < 1e-10最大 1.24e-14通过
能训练3 轮 MNIST98.52%通过
能解释慢说出瓶颈Python 位置循环 746.9 ns通过

用一段断言把前向对齐钉死:

python
# 第 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 时这个数会变成多少?