img2col 卷积优化讲解

转载：https://juejin.cn/post/7068113084451127333

目录：

简介
原理
Input Features -> Input Matrix
Convolution Kernel -> Kernel Matrix
Input Matrix * Kernel Matrix = Output Matrix
结语
代码实现

简介

img2col 是一种实现卷积操作的加速计算策略。它能将卷积操作转化为 GEMM，从而最大化地缩短卷积计算的时间。

GEMM 是通用矩阵乘 (General Matrix Multiply) 的英文缩写，其实就是一般意义上的矩阵乘法，数学表达就是 C = A x B。根据上下文语境，GEMM 有时也指实现矩阵乘法的函数接口。

为什么要将卷积操作转化为 GEMM 呢？

• 因为线性代数领域已经有非常成熟的计算接口（BLAS，Fortran 语言实现）来高效地实现大型的矩阵乘法，几乎可以做到极限优化。
• 将卷积过程中用到的所有特征子矩阵整合成一个大型矩阵存放在连续的内存中，虽然增加了存储成本，但是减少了内存访问的次数，从而缩短了计算时间。原理img2col 的原理可以用下面这一张图来概括：
  

  在这里插入图片描述

Input Features -> Input Matrix

不难看出，输入特征图一共有三个通道，我们以不同的颜色来区分。

在这里插入图片描述

以蓝色的特征图为例，它是一个 3 x 3 的矩阵，而卷积核是一个 2 x 2 的矩阵，当卷积核的滑动步长为 1 时，那么传统的直接卷积计算一共需要进行 4 次卷积核与对应特征子矩阵之间的点积运算。

现在我们把每一个特征子矩阵都排列成一个行向量（如图中编号1️⃣、2️⃣所示），然后把这 4 个行向量堆叠成一个新的矩阵，就得到了蓝色特征图所对应的 Input Matrix。

当输入特征图不止一个通道时，则对每一个通道的特征图都采用上述操作，然后再把每一个通道对应的 Input Matrix 堆叠成一个完整的 Input Matrix。

Convolution Kernel -> Kernel Matrix

不难看出，卷积核一共有两个，每个均为三通道，我们以第一个卷积核为例进行讲解。

在这里插入图片描述

将卷积核转化成矩阵的方式和第一步有些类似，只是这里应该转化成列向量（如图中编号1️⃣、2️⃣、3️⃣所示）。如果第一步转化成列向量，则这里应该转化成行向量，这是由矩阵乘法的计算特性决定的，即一个矩阵的每一行和另一个矩阵的每一列做内积，所以特征图和卷积核只能一个展开为行，一个展开为列。

同样地，如果卷积核有多个通道，则对每一个通道的卷积核都采用上述操作，然后再把每一个通道对应的 Kernel Matrix 堆叠成一个完整的 Kernel Matrix。

Input Matrix * Kernel Matrix = Output Matrix

在得到上述两个矩阵之后，接下来调用 GEMM 函数接口进行矩阵乘法运算即可得到输出矩阵，然后将输出矩阵通过 col2img 函数就可以得到和卷积运算一样的输出特征图。

在这里插入图片描述

结语

通过 img2col 函数，我们只需执行一次矩阵乘法计算就能得到与卷积运算相同的结果，而传统的直接卷积计算光是一个通道就需要进行 4 次（仅指本例中）卷积核与对应特征子矩阵之间的点积运算，那么如果通道数特别多？输入特征图非常庞大呢？那计算的次数将是成倍增长的！

有些同学可能会担心将所有特征子矩阵都堆叠到一个矩阵中，会不会导致内存不够用或者计算速度非常慢，尤其是在深度神经网络中。其实不用担心，因为矩阵的存储和计算其实都是非常规则的，很容易通过分布式和并行的方式来解决，感兴趣的同学可以自行阅读相关论文。

代码实现

https://github.com/lixiang007666/idealflow/blob/a5f80d68f60411607f65291c4472bf94389b6647/idealflow/core/utils.py#L6