卷积怎么实现?手写 CNN 才让我真正搞懂 im2col

7 阅读3分钟

今天上传了新的章节:卷积神经网络(6)卷积层。这是到目前为止计算编程难度最大的一个章节,值得总结一下。


卷积的概念很好理解:一个小卷积核在图像上滑动,每次对覆盖区域做加权求和。

这个"滑动"用代码怎么实现?

最朴素的写法是几层嵌套循环:遍历每张图,遍历每个输出位置,遍历卷积核的每个元素...

光是看这几层循环,性能就已经没救了。

PyTorch 帮我们把这一切都藏起来了,nn.Conv2d 一行搞定。直到开始动手手搓 CNN,才不得不去探究背后的实现技巧。

主流的解决方案就是 im2col(图转列)

im2col 的核心思想

im2col 的思路非常直接:把卷积操作变成矩阵乘法

卷积的每一个输出值,都是"输入图像的某个局部区块"与"卷积核"的内积。一张图像有多少个输出位置,就有多少次这样的内积。

im2col 做的事情是:

  1. 把每个输出位置对应的输入区块展平成一行
  2. 所有输出位置的行拼在一起,构成一个大矩阵;
  3. 把卷积核也展平成一个向量;
  4. 整个卷积就变成了一次矩阵乘法

矩阵乘法是 NumPy(以及所有深度学习框架底层)最擅长的操作,速度远超多层循环。

一个具体的例子

假设输入数据是一张 1×4×4 的图像(通道:1;高:4;宽:4),使用 3 个卷积核是 3×2×2(输出通道:3;高:2;宽:2),步长为 1。

输入图像(2×1×4×4):     
1  2  3  4              
5  6  7  8              
9  10 11 12  
13 14 15 16 

3 个卷积核(3×2×2):
1  0  |  1  1  |  0  1 
0  1  |  0  0  |  0  1 

输出尺寸是 3×3×3 (通道:3;高:3;宽:3)。

卷积核是 3 个输出通道,每次输出 3 个位置,一共需要 9 次加权求和。

im2col 把这 18 次计算对应的输入区块各展平成一行:

位置(0,0) → [1,  2,  5,  6 ] 
位置(0,1) → [2,  3,  6,  7 ] 
位置(0,2) → [3,  4,  7,  8 ] 
位置(1,0) → [5,  6,  9,  10]
位置(1,1) → [6,  7,  10, 11] 
位置(1,2) → [7,  8,  11, 12] 
位置(2,0) → [9,  10, 13, 14] 
位置(2,1) → [10, 11, 14, 15] 
位置(2,2) → [11, 12, 15, 16] 

拼在一起就是一个 9×4 的矩阵(输出尺寸:9;核尺寸:4)。

卷积核展平成 3×4 的向量:

[1, 0, 0, 1] 
[1, 1, 0, 0] 
[0, 1, 0, 1]

做一次矩阵乘法,直接得到 9×3 个输出值,再 reshape 成 3×3×3 就完成了。

关键实现:stride_tricks

im2col 的朴素写法是把那些区块逐个 copy 出来,但这会产生大量内存复制。

NumPy 有一个更巧妙的做法:np.lib.stride_tricks.as_strided

它不复制数据,而是通过调整内存步长(stride)来创建一个新的"视图" ,让我们可以用不同的下标规则访问同一块内存。


前向传播有了,但训练还需要把梯度传回去。

im2col 的反向就是  "col2im" :把输出位置的梯度分散回输入的对应区块。

每个输入像素可能被多个输出位置用到(相邻的卷积窗口会重叠),所以要用 np.add.at 做累加而不是赋值。


im2col 本质上是一种以空间换时间的策略:

  • 代价:im2col 矩阵比原始输入大。对于 k×k 的卷积核,矩阵大约膨胀 k² 倍;
  • 收益:所有计算都变成矩阵乘法,NumPy 的向量化能力被充分利用,比嵌套循环快几十倍甚至更多。

这也是为什么 PyTorch、TensorFlow 在 CPU 上的卷积底层都在用类似的技术(GEMM-based convolution)。


📖 完整章节(可运行 Jupyter Notebook):n2gpt.github.io/from-neuron…

整个系列从单个神经元出发,纯 NumPy,一路实现到 GPT。欢迎 Star、欢迎提问。