深入解析AlexNet卷积公式：从数学原理到高效实现

1次阅读

共计 1990 个字符，预计需要花费 5 分钟才能阅读完成。

2012 年 ImageNet 竞赛中，AlexNet 以绝对优势夺冠，标志着深度学习在计算机视觉领域的崛起。其核心创新之一就是高效利用卷积运算提取空间特征。相比全连接层，卷积层的参数共享特性大幅减少了模型参数量，使得训练深层网络成为可能。

AlexNet 中标准卷积运算可表示为：

$$O_{i,j,k} = \sum_{m=1}^{C_{in}} \sum_{u=-k}^{k} \sum_{v=-k}^{k} W_{u,v,m,k} \cdot I_{i\cdot s+u+p, j\cdot s+v+p, m} + b_k$$

其中关键参数说明：

$C_{in}/C_{out}$: 输入 / 输出通道数
$k$: 卷积核半径（实际尺寸为 $(2k+1)\times(2k+1)$）
$s$: 步长(stride)
$p$: 填充(padding)
$W$: 卷积核权重张量

输出特征图尺寸计算公式：

$$H_{out} = \left\lfloor \frac{H_{in} + 2p – (2k+1)}{s} \right\rfloor + 1$$

通过复杂度分析可发现痛点：

计算复杂度：$O(H_{out} \times W_{out} \times C_{out} \times (2k+1)^2 \times C_{in})$
内存访问：每个输入元素被重复读取 $(2k+1)^2 \times C_{out}$ 次
并行度利用：传统循环实现难以充分利用 GPU 多核特性

实测 AlexNet 第一个卷积层（输入 224x224x3，核 11×11，stride=4）：

理论 FLOPs：约 704M 次运算
显存访问量：超过 5GB

将卷积运算转换为矩阵乘法：

通过 im2col 操作将输入展开为 $(H_{out}W_{out}) \times (C_{in}(2k+1)^2)$ 矩阵
卷积核重排为 $(C_{in}(2k+1)^2) \times C_{out}$ 矩阵
调用 cuBLAS 的 SGEMM 进行计算

PyTorch 实现核心代码：

def conv_im2col(x, weight, bias, stride, padding):
    # 输入 x: [N,C_in,H,W]
    N, C_in, H_in, W_in = x.shape
    C_out, _, k, _ = weight.shape

    # 计算输出尺寸
    H_out = (H_in + 2*padding - k) // stride + 1
    W_out = (W_in + 2*padding - k) // stride + 1

    # im2col 转换
    cols = F.unfold(x, kernel_size=k, padding=padding, stride=stride)  # [N, C_in*k*k, H_out*W_out]
    cols = cols.permute(0, 2, 1).reshape(N*H_out*W_out, -1)

    # 矩阵乘法
    weight_flat = weight.view(C_out, -1).t()  # [C_in*k*k, C_out]
    out = torch.mm(cols, weight_flat) + bias

    return out.view(N, H_out, W_out, C_out).permute(0, 3, 1, 2)

直接编写 CUDA 核函数实现：