从零手写卷积与池化:im2col原理、反向传播与梯度校验实战
1. 从全连接到卷积为什么我们要自己造这个轮子如果你跟着这个系列一路走到现在说明你已经手搓过一个能跑通全连接网络的小型深度学习框架了。反向传播的链式法则、张量的基本运算、优化器的参数更新这些骨架都已经立起来了。但当你拿着这个框架去碰真实图像数据的时候大概率会撞上一堵墙——全连接层处理图像参数量会爆炸。我拿一个最朴素的场景算笔账。一张 224×224 的彩色图片展平之后是 224×224×3 150528 个输入维度。如果第一层隐藏层只放 1000 个神经元光这一层的权重矩阵就是 150528×1000 ≈ 1.5 亿个参数。这还只是第一层。更要命的是全连接层把图像的空间结构彻底打碎了左上角的像素和右下角的像素被同等对待局部特征、平移不变性这些图像里最本质的东西全丢了。卷积层和池化层就是来解决这两个问题的。卷积通过局部连接和权重共享把参数量压下来池化通过下采样进一步压缩空间维度并带来一定的平移鲁棒性。这一篇我就把这两个层从零实现一遍包括前向传播和反向传播的完整推导与代码落地。适合已经理解反向传播基本原理、想搞清楚卷积底层计算细节的人。如果你只是会调 PyTorch 的nn.Conv2d那这篇正好帮你把那层黑盒拆开。2. 卷积层与池化层的整体设计思路2.1 核心需求拆解我们到底要实现什么在动手写代码之前得先把需求理清楚。一个能用的卷积层至少要满足这几个条件支持多通道输入和多通道输出、支持可配置的卷积核大小、步长和填充、前向传播能算出正确的输出特征图、反向传播能正确算出对输入、对权重、对偏置的梯度。池化层相对简单但要支持最大池化和平均池化两种模式同样要能反向传播。这里有个设计上的关键决策卷积怎么实现。常见的有两种路子一种是im2colimage to column把卷积转成矩阵乘法另一种是直接多层循环暴力计算。我在这个框架里选择了im2col原因很直接——矩阵乘法是数值计算里优化最成熟的操作一旦转成矩阵乘后面接 BLAS 库或者自己写的优化矩阵乘都能直接复用性能比六层嵌套循环高出一个量级。代价是要多占一点内存因为输入会被展开成冗余的列矩阵但对于教学和中小规模实验来说这个取舍完全值得。2.2 张量的内存布局约定在写任何代码之前必须先定死张量的维度顺序否则后面全是坑。我采用的是NCHW布局也就是(batch, channel, height, width)。为什么选 NCHW 而不是 NHWC因为在 CPU 上用im2col做矩阵乘的时候NCHW 展开出来的列矩阵在通道维度上是连续的访存局部性更好。而且这个布局和主流框架的默认约定一致后面做对比验证的时候不容易搞混。权重张量的形状我定为(out_channels, in_channels, kernel_h, kernel_w)。这个顺序很重要out_channels在最前面意味着每个输出通道对应一个独立的卷积核组做矩阵乘的时候可以直接把权重 reshape 成(out_channels, in_channels*kernel_h*kernel_w)和im2col展开后的列矩阵直接相乘。2.3 前向与反向的数学骨架卷积的前向传播本质上是一个互相关运算严格数学意义上的卷积会翻转核但深度学习里大家约定俗成都用互相关PyTorch 也是这么做的。输出特征图上某个位置的值等于输入对应感受野和卷积核逐元素相乘再求和加上偏置。反向传播要算三个梯度对输入的梯度、对权重的梯度、对偏置的梯度。偏置的梯度最好算就是把输出梯度在 batch 和空间维度上求和。权重的梯度可以通过im2col后的输入矩阵和输出梯度的矩阵乘得到。输入的梯度稍微绕一点需要把输出梯度做一次反向的 im2col也就是col2im把列矩阵还原回输入的形状并累加。池化层的反向传播有个容易踩的坑最大池化在前向时记录了最大值的位置反向时梯度只回传到那个位置其他位置是零。平均池化则是把梯度平均分配到整个池化窗口。这个差异必须在实现时严格区分。3. 核心细节解析与实操要点3.1 im2col 的原理与实现细节im2col这个名字直译就是图像转列。它的核心思想是把输入特征图上每一个会被卷积核覆盖的感受野拉成一个列向量所有感受野的列向量拼在一起就形成了一个大矩阵。假设输入是(N, C, H, W)卷积核KH×KW步长S填充P那么输出空间尺寸是out_h (H 2P - KH) / S 1 out_w (W 2P - KW) / S 1展开后的列矩阵形状是(N * out_h * out_w, C * KH * KW)。每一行对应一个输出位置每一列对应卷积核里的一个权重位置。这样卷积就变成了权重矩阵 (out_channels, C*KH*KW)乘以列矩阵的转置得到(out_channels, N*out_h*out_w)再 reshape 回(N, out_channels, out_h, out_w)。实现的时候有个细节要注意填充操作我建议在im2col内部处理而不是先对输入做一次物理填充。因为物理填充会复制一份完整数据内存开销大。在im2col里判断索引是否越界越界就填零这样更省内存。但代价是每次取值都要做边界判断速度会慢一点。我实测下来对于中小规模输入这个开销可以接受。def im2col(x, kernel_h, kernel_w, stride, pad): N, C, H, W x.shape out_h (H 2 * pad - kernel_h) // stride 1 out_w (W 2 * pad - kernel_w) // stride 1 # 先做物理填充简化索引逻辑 x_padded np.pad(x, ((0,0),(0,0),(pad,pad),(pad,pad)), modeconstant) col np.zeros((N, C, kernel_h, kernel_w, out_h, out_w)) for y in range(kernel_h): y_max y stride * out_h for x_ in range(kernel_w): x_max x_ stride * out_w col[:, :, y, x_, :, :] x_padded[:, :, y:y_max:stride, x_:x_max:stride] col col.transpose(0, 4, 5, 1, 2, 3).reshape(N * out_h * out_w, -1) return col这段代码用的是按核位置切片的写法比六层循环快很多。它的思路是对卷积核的每一个位置(y, x_)一次性把所有 batch、所有通道、所有输出位置的对应输入值取出来。这种向量化写法是性能的关键。3.2 卷积层前向传播的完整实现有了im2col前向传播就清爽了。权重 reshape 成二维和列矩阵做矩阵乘加上偏置再 reshape 回 NCHW。class Conv2D: def __init__(self, in_channels, out_channels, kernel_size, stride1, pad0): self.W np.random.randn(out_channels, in_channels, kernel_size, kernel_size) * 0.01 self.b np.zeros(out_channels) self.stride stride self.pad pad self.kernel_size kernel_size def forward(self, x): self.x_shape x.shape N, C, H, W x.shape KH KW self.kernel_size out_h (H 2*self.pad - KH) // self.stride 1 out_w (W 2*self.pad - KW) // self.stride 1 self.col im2col(x, KH, KW, self.stride, self.pad) self.col_W self.W.reshape(self.W.shape[0], -1) out self.col self.col_W.T self.b out out.reshape(N, out_h, out_w, -1).transpose(0, 3, 1, 2) return out注意最后那个transpose(0, 3, 1, 2)因为矩阵乘出来的结果是(N*out_h*out_w, out_channels)需要先 reshape 成(N, out_h, out_w, out_channels)再转成(N, out_channels, out_h, out_w)。这个维度变换的顺序很容易写错我第一次实现的时候就因为把out_h和out_channels搞反了导致输出图像完全错乱排查了半天。3.3 反向传播col2im 与梯度计算反向传播是这一篇最难的部分。先看权重梯度dW dout_reshaped.T col其中dout_reshaped是输出梯度 reshape 成(N*out_h*out_w, out_channels)。偏置梯度db dout.sum(axis(0,2,3))。输入梯度的计算需要col2im也就是把列矩阵还原回输入形状。这里的关键是输入的一个像素可能被多个感受野覆盖所以梯度要累加而不是覆盖。def col2im(col, x_shape, kernel_h, kernel_w, stride, pad): N, C, H, W x_shape H_padded, W_padded H 2*pad, W 2*pad x_padded np.zeros((N, C, H_padded, W_padded)) col col.reshape(N, -1, C, kernel_h, kernel_w).transpose(0, 2, 3, 4, 1) out_h (H 2*pad - kernel_h) // stride 1 out_w (W 2*pad - kernel_w) // stride 1 for y in range(kernel_h): y_max y stride * out_h for x_ in range(kernel_w): x_max x_ stride * out_w x_padded[:, :, y:y_max:stride, x_:x_max:stride] col[:, :, y, x_, :, :] return x_padded[:, :, pad:Hpad, pad:Wpad]注意这里是而不是这是整个反向传播正确性的命门。如果写成赋值重叠区域的梯度就丢了训练出来的网络会莫名其妙地不收敛。3.4 池化层的实现要点池化层比卷积层简单但最大池化的反向传播有个经典陷阱。前向的时候必须记录每个池化窗口里最大值的位置反向的时候才能把梯度准确地回传。我见过不少人图省事反向时重新比较一遍输入找最大值位置如果输入在前向和反向之间被修改过比如某些 in-place 操作就会出错。稳妥的做法是前向时就把 argmax 存下来。class MaxPool2D: def __init__(self, pool_size2, stride2): self.pool_size pool_size self.stride stride def forward(self, x): N, C, H, W x.shape PH PW self.pool_size out_h (H - PH) // self.stride 1 out_w (W - PW) // self.stride 1 self.x x self.arg_max np.zeros((N, C, out_h, out_w), dtypenp.int32) out np.zeros((N, C, out_h, out_w)) for i in range(out_h): for j in range(out_w): window x[:, :, i*self.stride:i*self.stridePH, j*self.stride:j*self.stridePW] out[:, :, i, j] window.max(axis(2,3)) self.arg_max[:, :, i, j] window.reshape(N, C, -1).argmax(axis2) return out反向的时候根据arg_max把梯度散射回原来的位置。平均池化则简单得多梯度平均分配到窗口内每个位置不需要记录位置。4. 实操过程与核心环节实现4.1 搭建一个可训练的迷你 CNN光有层还不够得把它们串起来跑通一个完整的训练流程才能验证实现是否正确。我搭了一个结构很简单的网络Conv2D(1, 8, 3, stride1, pad1) - ReLU - MaxPool2D(2, 2) - Conv2D(8, 16, 3, stride1, pad1) - ReLU - MaxPool2D(2, 2) - Flatten - Linear - Softmax。输入是 28×28 的单通道图像输出 10 类。这里有个维度衔接的细节要算清楚。28×28 经过第一个卷积pad1kernel3stride1输出还是 28×28池化后变 14×14。第二个卷积输出 14×14池化后变 7×7。所以 Flatten 之后的维度是 16×7×7 784。这个数字必须和全连接层的输入维度对上否则矩阵乘直接报错。4.2 梯度数值校验别信你的直觉写完反向传播千万别急着训练先做数值梯度校验。这是我在这个项目里踩过最大的坑之后养成的习惯。方法很简单用数值微分算一个近似的梯度和你反向传播算出来的解析梯度对比相对误差在 1e-6 量级才算通过。def numerical_gradient(f, x, eps1e-4): grad np.zeros_like(x) it np.nditer(x, flags[multi_index]) while not it.finished: idx it.multi_index old x[idx] x[idx] old eps fx1 f(x) x[idx] old - eps fx2 f(x) grad[idx] (fx1 - fx2) / (2 * eps) x[idx] old it.iternext() return grad校验的时候有个技巧把网络里的 ReLU 暂时换成恒等映射或者确保输入不会落在 ReLU 的不可导点上。因为 ReLU 在零点不可导数值校验会在这里产生巨大误差让你误以为反向传播写错了。我第一次校验时就被这个坑了整整一个下午后来把激活函数换成 Sigmoid 才顺利通过。4.3 训练过程与损失曲线观察校验通过之后就可以训练了。我用的是一个手写数字数据集batch size 设 32学习率 0.01用 SGD 优化器。前几个 epoch 损失下降得很明显大概 5 个 epoch 之后准确率能到 95% 以上。如果损失不降反升或者震荡得厉害通常是这几个原因学习率太大、权重初始化尺度不对、或者反向传播里某个梯度符号写反了。我建议在训练循环里每隔一定步数打印一次损失并且把损失曲线画出来。如果曲线是锯齿状剧烈震荡先把学习率砍一半试试。如果损失一直卡在一个较高的值不动检查一下im2col和col2im的索引是否对称这两个函数只要有一个索引偏了梯度就是错的网络根本学不动。4.4 性能优化从能跑到跑得快最初的实现跑一个 epoch 要几分钟慢得让人抓狂。我做了几处优化速度提升了大概三倍。第一处是把im2col里的循环尽量向量化用切片代替逐元素赋值。第二处是确保矩阵乘用的是优化过的库函数而不是自己写的三重循环。第三处是减少不必要的内存拷贝比如col2im里直接操作填充后的数组最后再切掉填充部分。还有一个容易被忽略的点数据类型。默认的 float64 精度高但慢改成 float32 之后速度明显提升而且对训练精度几乎没有影响。这个改动一行代码的事收益却很大。5. 常见问题与排查技巧实录5.1 维度不匹配问题速查维度错误是手写卷积层最高频的问题。我把常见的错误现象和原因整理成了一张表方便对照排查。错误现象可能原因排查方法矩阵乘报维度错误im2col 输出列数不等于 C×KH×KW打印 col.shape 和 col_W.shape 对比输出特征图尺寸不对out_h/out_w 公式算错手算一遍 (H2P-KH)/S1reshape 后图像错乱transpose 维度顺序写错用小尺寸输入逐步打印 shape反向传播梯度形状对不上col2im 的 reshape 顺序和前向不一致检查 transpose 参数是否对称5.2 梯度消失与爆炸的应对卷积层堆深了之后梯度消失或爆炸几乎必然出现。我在实验里发现权重初始化用小的随机数比如乘 0.01能缓解爆炸但层数一多还是会有消失问题。几个实用的应对手段在卷积层后面接 BatchNorm这个后续可以单独实现、用 ReLU 而不是 Sigmoid 作为激活函数、控制网络深度。如果训练时损失突然变成 NaN八成是梯度爆炸了先把学习率降一个数量级再检查权重初始化。5.3 池化层反向传播的隐蔽 bug最大池化反向传播有个很隐蔽的 bug当池化窗口里有多个相同的最大值时argmax只会返回第一个位置梯度只回传到那一个位置。这在数学上其实没问题因为梯度回传到任意一个最大值位置都等价但如果你的实现里用了某种平均分配的逻辑就会出错。我建议严格按argmax的位置回传不要自作聪明。提示调试池化层反向传播时可以构造一个输入让池化窗口里只有一个明显的大值其余都是小值这样梯度应该只出现在那一个位置。如果发现梯度散布在多个位置说明实现有问题。5.4 内存占用过大的处理im2col的代价是内存。一个(32, 3, 224, 224)的输入经过 3×3 卷积展开后列矩阵大小是32×224×224 × 27约 4300 万个浮点数float32 下就是 170MB。如果 batch 再大一点内存直接爆掉。我的处理办法是训练时用较小的 batch size或者把im2col分块进行每次只展开一部分输出位置。对于教学框架控制 batch size 是最简单有效的办法。6. 我个人在实际操作中的几点体会把卷积层和池化层手写一遍最大的收获不是代码本身而是对为什么卷积能work这件事有了肌肉记忆般的理解。以前调nn.Conv2d的时候stride、padding、kernel_size这些参数就是几个数字现在我能清楚地知道每个参数变化时特征图的尺寸怎么变、感受野怎么变、参数量怎么变。如果让我给正在跟着实现的朋友一个建议那就是先让数值梯度校验通过再谈训练。我见过太多人反向传播写错了然后花大量时间去调学习率、换初始化其实问题根本不在那里。数值校验虽然慢但它是唯一能告诉你梯度算对了没有的可靠手段。另外im2col这个技巧值得单独记牢。它不只用在卷积里很多需要滑动窗口的操作都能用它转成矩阵乘。理解了它你再看一些高性能推理框架的实现会发现到处都是它的影子。这个系列走到这里框架已经能处理图像了下一步自然是把 BatchNorm、Dropout 这些训练技巧补上让网络真正具备实战能力。