vLLM 卷积计算加速｜img2col 原理详解

Table of Contents

一、引言
#

卷积运算是 VL 等多模态模型在处理图像、视频等数据时的核心步骤之一，使用 img2col 算法可以将输入数据和卷积核展平为两个大的矩阵，并通过一次高效的矩阵乘法得到卷积结果，从而极大地提升计算的效率。本文将详细讲解 img2col 算法的基本原理和代码实现，并对 vLLM 中的卷积算子进行介绍。

二、卷积的基本原理
#

在讲解 img2col 算法之前，我们先简单介绍下什么是卷积运算（Convolution）。

简单来说，卷积运算就是一个小窗口（一般称为“卷积核”或“滤波器”）在一个大的输入数据（如图片）上滑动，并在每个位置计算点积，最终生成一个新的、更精炼的特征图的过程。其中，卷积核一般使用正方形，比如在下图中，使用的就是一个 3 x 3 的卷积核（卷积核的通道数 = 输入的通道数，一般为 3，代表图像的红、绿、蓝三钟颜色）。

卷积的计算过程：

属于同一输入通道的卷积核在对应的图像数据上进行滑动，并在每一个位置处计算这 9 个数据的点积和；
将每个输入通道的计算结果在每个位置上进行相加，得到形状为 (1, 3, 3) 的输出；
当有多个卷积核（卷积核的个数等于输出通道数）时，将输入数据依次和每个卷积核进行步骤 1. 和 2. 的计算，最终将所有结果拼接为一个形状为 (m, 3, 3) 张量。

注意：上图摘自“刘二大人”的 PyTorch 深度学习实践，感兴趣的朋友可以自行移步 B 站进行学习；关于卷积更多参数的含义和用法，也可以自行浏览 PyTorch 官方文档进行了解。

三、img2col 原理详解
#

在了解了卷积的基本原理之后，下面我们将通过一个简单的例子来对 img2col 的原理进行讲解。

老规矩，先上图，一切尽在图中～

img2col 的计算过程：

将输入的图像数据展平，并按 kernel_size 进行分块（这里展示的是 kernel_size = stride 步长的特殊情况）；
交换张量维度，外层是 (2, 2) 的分块，内层是划分后 (3, 3, 3) 的输入数据；
将每一个 (3, 3, 3) 的数据块展平为一行（27 列），共 4 行；
将每一个卷积核展开为一列（27 行），共 2 列（两个卷积核，out_channels = 2）；
两个大矩阵直接做一把 matmul（可以用到高度优化的 GEMM 库，计算效率高）；
将矩阵运算的结果 reshape 为标准卷积的输出形式。

高清图片链接：link，画图不易，走过路过欢迎点一个 Star！

PyTorch 代码实现：

def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor:
    assert x.dim() == 4
    B, C, H, W = x.shape
    K1, K2 = self.kernel_size
    H, W = H // K1, W // K2
    x = x.unfold(2, K1, K1).unfold(3, K2, K2)
    x = x.permute(0, 2, 3, 1, 4, 5).reshape(-1, self.input_size)
    x = F.linear(
        x,
        self.weight.view(self.out_channels, self.input_size),
        self.bias,
    )
    x = x.view(B, H, W, self.out_channels).permute(0, 3, 1, 2)
    return x

需要特别注意的是：只有当 kernel_size = stride 时才适合使用上述优化，因为此时卷积核在每次移动之后，处理的数据互相不重叠，能够完美地展开为一个大的矩阵。当不满足上述条件时，重叠的窗口会导致重排矩阵中存在大量的数据冗余，从而带来更大的内存占用。特别是对于大尺寸的输入、大卷积核或小步长，这个重排矩阵会非常庞大，可能成为内存的瓶颈。