Loading... 笔者在工作的过程中,使用Triton和cuda实现或优化过rwkv6的内核,中途也遇到过一些数值精度的陷阱。于是提起用Triton实现深度学习常见的op的心思,虽然有一些实现的计划,从简单到复杂一步步入门常见kernel的实现,但是正好朋友在FPGA上面实现卷积算法,于是今天就带着读者了解一下什么是卷积,如何高效率的实现卷积,当然,如果笔者数学功底还在的话,会带着实现卷积的反向传播。 ## 什么是卷积 笔者假设读者具有一定的高等数学功底,倘若读者数学不好也没有关系,笔者会在每个算法提供代码和数学公式来帮助大家入门深度学习的数学世界。(笑,因为笔者在本科上学的时候深受数学的折磨,后面继续深造和工作的时候发现对笔者而言,最好的理解数学的方法是写一遍代码来理解数学公式背后的逻辑和作用。 > 在[泛函分析](https://zh.wikipedia.org/wiki/%E6%B3%9B%E5%87%BD%E5%88%86%E6%9E%90 "泛函分析")中,**卷积**(convolution),或译为**叠积**、**褶积**或**旋积**,是透过两个[函数](https://zh.wikipedia.org/wiki/%E5%87%BD%E6%95%B0 "函数")$f$和$g$生成第三个函数的一种数学[算子](https://zh.wikipedia.org/wiki/%E7%AE%97%E5%AD%90 "算子"),表征函数$f$与经过翻转和平移的$g$ 卷积是数学分析中一种重要的运算。设 $f(t)$ 和 $g(t)$ 是实数取上的两个可积函数,定义二者的卷积 $(f * g)(t)$ 为如下特定形式的积分变换: $$ (f * g)(t) \triangleq \int_{-\infty}^{\infty} f(\tau)g(t - \tau) d\tau $$ $(f * g)(t)$ 仍为可积函数,并且有: $$ (f * g)(t) \triangleq \int_{-\infty}^{\infty} f(t - \tau)g(\tau) d\tau = (g * f)(t) $$ 函数 $f$ 和 $g$, 如果只支撑在 $[0, \infty]$ 上,则积分界限可以截断为: $$ (f * g)(t) = \int_0^t f(\tau)g(t - \tau) d\tau \quad 对于 f,g : [0,\infty) \to \mathbb{R} $$ 对于两个值出复数值的多元实变函数,可以定义二者的卷积为如下形式的多重积分: $$ (f * g)(t_1, t_2, \cdots, t_n) \triangleq \int \int \cdots \int_{\mathbb{R}^n} f(\tau_1, \tau_2, \cdots, \tau_n)g(t_1 - \tau_1, t_2 - \tau_2, \cdots, t_n - \tau_n) d\tau_1 d\tau_2 \cdots d\tau_n $$ $$ \triangleq \int_{\mathbb{R}^n} f(\tau)g(t - \tau) d^n \tau $$ 卷积有个通用的工程上的符号定义: $$ f(t) * g(t) \triangleq \int_{-\infty}^{\infty} f(\tau)g(t - \tau) d\tau $$ 它必须被谨慎解释以避免混淆。例如: $f(t) * g(t - t_0)$ 等价于 $(f * g)(t - t_0)$, 而 $f(t - t_0) * g(t - t_0)$ 则实际上等价于 $(f * g)(t - 2t_0)$。 (以上是维基百科的解释,我猜作为初学者的你一定在公式里迷失了,没关系。我们可以举个例子。 例1: 矩形函数的卷积 假设我们有两个简单的矩形函数: $$ f(t) = \begin{cases} 1, & 0 \leq t \leq 1 \\ 0, & \text{其他} \end{cases} $$ $$ g(t) = \begin{cases} 1, & 0 \leq t \leq 2 \\ 0, & \text{其他} \end{cases} $$ 它们的卷积 $(f * g)(t)$ 可以直观地理解为: $$ (f * g)(t) = \begin{cases} t, & 0 \leq t < 1 \\ 1, & 1 \leq t < 2 \\ 3-t, & 2 \leq t \leq 3 \\ 0, & \text{其他} \end{cases} $$ 这个结果形状像一个梯形。 例2: 指数函数与自身的卷积 考虑函数 $f(t) = e^{-at}$ 其中 $a > 0$ 且 $t \geq 0$。它与自身的卷积为: $$ (f * f)(t) = \int_0^t e^{-a\tau} e^{-a(t-\tau)} d\tau = te^{-at} $$ 这个结果显示,两个指数函数的卷积产生了一个新的函数,它在开始时增加,然后逐渐减小。 例3: 高斯函数的卷积 高斯函数(正态分布)的卷积特别有趣。设: $$ f(t) = \frac{1}{\sigma\sqrt{2\pi}} e^{-\frac{t^2}{2\sigma^2}} $$ $$ g(t) = \frac{1}{\tau\sqrt{2\pi}} e^{-\frac{t^2}{2\tau^2}} $$ 它们的卷积结果仍然是一个高斯函数: $$ (f * g)(t) = \frac{1}{\sqrt{2\pi(\sigma^2 + \tau^2)}} e^{-\frac{t^2}{2(\sigma^2 + \tau^2)}} $$ 这个性质在概率论和信号处理中非常有用。 笔者在这里引入了几个高等数学中的例子,帮助大家初步的了解数学原理。但值得大家庆幸的是,数学上的卷积积分范围通常是无限且连续的,但是在深度学习和实际工作中,卷积通常是在离散数据和有限的输入范围上处理的。所以上面的“入门介绍”大家可以通通忘掉:) ## 深度学习中的卷积 笔者在这里形象地定义深度学习中的卷积:深度学习中的卷积是一种使用有限大小的卷积核(filter)在有限大小图像输入上通过滑动进行数据处理的操作。在卷积操作中,卷积核按照一定的步长(stride)在输入图像上滑动,对于每个位置,将卷积核的参数与输入图像对应位置的局部区域数据逐个元素相乘,然后将这些乘积相加,得到一个标量值,作为输出特征图的一个元素。通过在整个输入图像上重复这个过程,我们得到一个新的输出特征图。 注:在实际应用中,输入通常是多通道的(如RGB图像),卷积核也会相应地有多个输出通道。此外,我们经常会使用填充(padding)来控制输出特征图的大小。这些因素共同构成了深度学习中的卷积操作。 ## 卷积的实际计算 我们还是用例子来说明,假设我们的输入图像是一个单通道5x5的图片,图片的数据如下: $$ \begin{bmatrix} 1 & 1 & 1 & 0 & 0 \\ 0 & 1 & 1 & 1 & 0 \\ 0 & 0 & 1 & 1 & 1 \\ 0 & 0 & 1 & 1 & 0 \\ 0 & 1 & 1 & 0 & 0 \end{bmatrix} $$ 我们有一个如下的卷积核: $$ \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 0 \\ 1 & 0 & 1 \end{bmatrix} $$ 那么通过计算我们可以看到,对于第一个计算的元素 $(1×1 + 0×1 + 1×1 + 0×0 + 1×1 + 0×1 + 1×0 + 0×0 + 1×1) = 4$ 重复这个过程,我们可以得到一个3x3的输出矩阵: $$ \begin{bmatrix} 4 & 3 & 4 \\ 2 & 4 & 3 \\ 2 & 3 & 4 \end{bmatrix} $$ 这个就是卷积的朴素算法,我们可以构建一个多重循环,来实现卷积的计算。 读者可以自行验证下面代码的正确性。 ```python import numpy as np def convolution(image, kernel): image_h, image_w = image.shape kernel_h, kernel_w = kernel.shape output_h = image_h - kernel_h + 1 output_w = image_w - kernel_w + 1 output = np.zeros((output_h, output_w)) for i in range(output_h): for j in range(output_w): output[i, j] = np.sum(image[i:i+kernel_h, j:j+kernel_w] * kernel) return output # 定义输入图像 image = np.array([ [1, 1, 1, 0, 0], [0, 1, 1, 1, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 0], [0, 1, 1, 0, 0] ]) # 定义卷积核 kernel = np.array([ [1, 0, 1], [0, 1, 0], [1, 0, 1] ]) # 执行卷积操作 result = convolution(image, kernel) print("卷积结果:") print(result) ``` ### img2col 算法 如果读者有线性代数的背景知识,那么我们可以有一种直觉,这种操作和矩阵乘矩阵有些类似。事实上我们确实有一种算法(被称作img2col, image to column),将卷积操作转换成矩阵的乘法操作。这种操作在很多情况下可以在GPU上更为高效的计算,因为它可以更好的分块和并行化。 > [矩阵 - 维基百科,自由的百科全书 (wikipedia.org)](https://zh.wikipedia.org/wiki/%E7%9F%A9%E9%98%B5) > [矩阵乘法-维基百科](https://zh.wikipedia.org/wiki/%E7%9F%A9%E9%98%B5%E4%B9%98%E6%B3%95) > [BLAS - 维基百科,自由的百科全书 (wikipedia.org)](https://zh.wikipedia.org/wiki/BLAS) 对于矩阵的乘法,我们知道,$A * B$是通过 A的行向量乘B的列向量并求和得到输出矩阵的每一个元素。这非常类似我们的卷积算法,那么问题便转化为:我们如何构建一个这样的A矩阵和B矩阵的问题。 我们先来考虑单通道图像的情况。 对于A矩阵,自然而然的,我们需要通过展平被卷积图像的局部区域来构建行向量。 对于B矩阵,同样的,我们可以通过展平卷积核参数来构建列向量。 最后,我们将img2col后的矩阵与展平的卷积核列向量相乘。 对于我们的例子,修正后的im2col矩阵应该是: $$ \begin{bmatrix} 1 & 1 & 1 & 0 & 1 & 0 & 1 & 0 & 1 \\ 1 & 1 & 0 & 1 & 1 & 1 & 0 & 0 & 0 \\ 1 & 0 & 0 & 1 & 1 & 1 & 1 & 1 & 0 \\ 0 & 1 & 1 & 0 & 1 & 1 & 0 & 1 & 1 \\ 1 & 1 & 1 & 1 & 1 & 0 & 0 & 1 & 0 \\ 1 & 1 & 0 & 1 & 0 & 0 & 1 & 1 & 0 \\ 0 & 1 & 0 & 0 & 1 & 1 & 1 & 0 & 0 \\ 1 & 1 & 1 & 1 & 1 & 0 & 1 & 0 & 0 \\ 1 & 0 & 0 & 1 & 0 & 0 & 1 & 0 & 0 \end{bmatrix} $$ 每一行代表输入图像中的一个3x3局部区域,展平为行向量。 展平的卷积核为列向量: $$ \begin{bmatrix} 1 \\ 0 \\ 1 \\ 0 \\ 1 \\ 0 \\ 1 \\ 0 \\ 1 \end{bmatrix} $$ 进行矩阵乘法后,我们得到: $$ \begin{bmatrix} 4 \\ 3 \\ 4 \\ 2 \\ 4 \\ 3 \\ 2 \\ 3 \\ 4 \end{bmatrix} $$ 这个结果向量包含了所有卷积操作的输出,将其重塑为3x3矩阵就得到了最终的卷积结果。 $$ \begin{bmatrix} 4 & 3 & 4 \\ 2 & 4 & 3 \\ 2 & 3 & 4 \end{bmatrix} $$ 同样的,笔者在这里不予解释的给出代码,希望读者可以自行实现: ```python import numpy as np def img2col(image, kernel_size): h, w = image.shape k_h, k_w = kernel_size # 计算输出大小 out_h = h - k_h + 1 out_w = w - k_w + 1 # 初始化输出矩阵 col = np.zeros((out_h * out_w, k_h * k_w)) # 填充col矩阵 for i in range(out_h): for j in range(out_w): patch = image[i:i+k_h, j:j+k_w].flatten() col[i*out_w + j] = patch return col def convolution_img2col(image, kernel): # 将kernel转换为列向量 k_flat = kernel.flatten().reshape(-1, 1) # 使用img2col将图像转换为矩阵 img_col = img2col(image, kernel.shape) # 进行矩阵乘法 result_flat = np.dot(img_col, k_flat) # 重塑结果为正确的输出形状 out_h = image.shape[0] - kernel.shape[0] + 1 out_w = image.shape[1] - kernel.shape[1] + 1 result = result_flat.reshape(out_h, out_w) return result # 定义输入图像 image = np.array([ [1, 1, 1, 0, 0], [0, 1, 1, 1, 0], [0, 0, 1, 1, 1], [0, 0, 1, 1, 0], [0, 1, 1, 0, 0] ]) # 定义卷积核 kernel = np.array([ [1, 0, 1], [0, 1, 0], [1, 0, 1] ]) # 执行img2col卷积操作 result = convolution_img2col(image, kernel) print("img2col卷积结果:") print(result) ``` 在这里笔者想略微考察一下读者,如果我们是一个多通道图像,这个问题又应该如何思考? 我们回顾线性代数的基础知识,一个 $m*n$的矩阵乘一个 $n*p$的矩阵,结果会是一个 $m*p$的矩阵。 也就是说,对于多通道图像的卷积操作,如果我们将其视为更大的矩阵乘法。假设我们有一个 3 通道的输入图像(如 RGB 图像),每个通道的大小仍为 5x5,我们的卷积核也相应地变为 3x3x3 的形状。 在这种情况下: 1. 输入图像可以被重塑为一个矩阵,其中每一行代表一个局部感受野,但现在每一行的长度是单通道情况的 3 倍。如果原来是 9 个元素(3x3),现在就是 27 个元素(3x3x3)。 2. 卷积核可以被展平为一个列向量,长度同样是 27。 3. 假设我们有 k 个这样的卷积核(也就是说,我们想要得到 k 个输出通道),那么我们可以将所有卷积核组合成一个 27 x k 的矩阵。 因此,我们的矩阵乘法变成了: (im2col 后的输入矩阵) × (展平的卷积核矩阵) = 输出矩阵 其中: - im2col 后的输入矩阵维度为 $(H'W') \times (C_{in}K_hK_w)$ - 展平的卷积核矩阵维度为 $(C_{in}K_hK_w) \times C_{out}$ - 输出矩阵维度为 $(H'W') \times C_{out}$ 这里: - $H'$ 和 $W'$ 是输出特征图的高度和宽度 - $C_{in}$ 是输入通道数 - $K_h$ 和 $K_w$ 是卷积核的高度和宽度 - $C_{out}$ 是输出通道数(即卷积核的数量) 读者可以自然的联想到,输入图像的通道不一定等于输出图像的通道,这正是深度学习中卷积层的工作原理,允许网络学习复杂的特征表示。 我们随后引出一种优化的img2col算法。读者可以做一个简单的计算,尽管展平后的矩阵和原始矩阵有着同样的数据,但是却扩大了矩阵的shape,这使得我们无论是在CPU上还是GPU上都使用了更多的内存,同时还引入了矩阵内存申请和分配的开销。 于是自然而然的,我们可以引入一种 Implicit GEMM(通用矩阵乘法,General Matrix Multiplication) 方法来计算卷积。 ### Implicit GEMM 读者可以带着问题来思考,为什么我们说通过GEMM的计算方法是更有优势的。 从上述算法的计算流程中可以看到,在img2col方法中,我们仅仅是通过重塑原始特征图的排布来构建新的矩阵,那么自然的,我们可以使用访问内存地址来索引原有的特征图,实现逻辑上的矩阵乘法。 首先,我们扩展完整的卷积操作,包括步长(stride),填充(padding)和膨胀(dilation)。 1. 步长:定义了卷积核每次操作在输入图像上滑动的距离,如果步长为1,那么卷积核每次都会移动一个pixel,如果我们有更大的步长,卷积核会跳过中间的一些像素。 2. 填充:读者可以自然的联想到,特征图不一定可以被卷积操作整数次处理,所以我们可以对边缘填充一些像素(通常我们会填充像素0),这样我们可以保持输出特征图的输出尺寸,同样的,这样可以缓解边缘像素参与计算次数较少的问题。 3. 膨胀:膨胀也被称作空洞卷积 膨胀率=1时,此时为标准的卷积操作,倘若膨胀率大于1,我们会在卷积操作的数据中间插入“空洞”。以下是一个简单的例子: 标准3x3卷积核: $$ \begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 & 6 \\ 7 & 8 & 9 \end{bmatrix} $$ 膨胀率为2的3x3卷积核: $$ \begin{bmatrix} 1 & 0 & 2 & 0 & 3 \\ 0 & 0 & 0 & 0 & 0 \\ 4 & 0 & 5 & 0 & 6 \\ 0 & 0 & 0 & 0 & 0 \\ 7 & 0 & 8 & 0 & 9 \end{bmatrix} $$ 当上述参数被引入卷积计算时,矩阵的索引往往会变得很复杂。 我们首先需要思考输入向量的维度和输出向量的维度,在此,由于矩阵操作天生的复杂性,笔者不得不引入一些字符来表征数据的维度。 笔者引入以下符号来表示数据的维度: - N: 批次大小 - C: 输入通道数 - H, W: 输入特征图的高度和宽度 - K: 输出通道数 (也是卷积核的数量) - R, S: 卷积核的高度和宽度 - U, V: 步长 (stride) 的高度和宽度方向 - pad_h, pad_w: 高度和宽度方向的填充 - dila_h, dila_w: 高度和宽度方向的膨胀率 根据这些参数,可以计算输出特征图的尺寸。(笔者建议可以先从5x5,8x8的输入手算一遍卷积,可以理解后续的公式) - P = (H + 2 * pad_h - dila_h * (R - 1) - 1) // U + 1 # 输出高度 - Q = (W + 2 * pad_w - dila_w * (S - 1) - 1) // V + 1 # 输出宽度 我们首先考虑没有任何padding,stride和dilation都是1的一种情况。在这种情况下,因为卷积核需要完全在输入内部的数据进行处理,所以行和列里多余的数据不会被处理。在这样的情况下,我们可以计算: $P=H-R+1$, 随后我们引入pad_h,有效高度等于$H+2*pad_h$, 所以我们现在有:$P = H + 2*padh - R+1$, 这两种情况是自然且易于推导的。 当引入步长和膨胀的时候,情况会变得略微复杂。 1. 引入步长。 1. 我们可以这样思考,没有引入步长的时候,输出的尺寸是$P = H + 2*padh - R+1$,引入步长,特征图的尺寸会变小。 2. 具体而言,输出尺寸是原有尺寸的 1/U, 由于卷积是在有效尺寸上输出的,所以我们需要向下取整,是用地板除法(floor division) 3. 最后我们得到: $P = (H + 2*padh - R+1) // U$ 2. 引入膨胀 1. 膨胀会扩大卷积核的有效尺寸,但不会增加参数数量。对于膨胀率dila_h,卷积核的有效高度变为:$R_{effective} = dila_h * (R - 1) + 1$ , 这是因为,我们从卷积核中心开始膨胀,那么R-1 是卷积核除中心外的高度,经过膨胀后,这部分的高度是:$dila_h * (R-1)$,最后我们加上中心的1个像素。 2. 这个有效高度代入之前的公式: $P = (H + 2*pad_h - R_{effective} + 1) // U$ $= (H + 2*pad_h - (dila_h * (R - 1) + 1) + 1) // U$ $= (H + 2*pad_h - dila_h * (R - 1) - 1) // U + 1$ 最后的 `+ 1` 确保了即使在极端情况下,输出高度至少为1 > 膨胀卷积,也称为空洞卷积(Atrous Convolution),是一种特殊的卷积操作,它在标准卷积的基础上引入了一个额外的参数:膨胀率(dilation rate)。这种卷积方式在保持参数数量和计算复杂度相对不变的情况下,能够显著增大感受野(receptive field)。 读者可以同样的去推导输出宽度的公式。 Implicit GEMM的核心思想是将卷积操作重塑为矩阵乘法。通过上述推导,我们定义如下的变量,其中重要的是和GEMM有关的参数,我们通过这些实现索引。 ```python N, C, H, W = x.shape K, C, R, S = w.shape U, V = stride pad_h, pad_w = padding dila_h, dila_w = dilation P = (H + 2 * pad_h - dila_h * (R - 1) - 1) // U + 1 # 输出高度 Q = (W + 2 * pad_w - dila_w * (S - 1) - 1) // V + 1 # 输出宽度 y = torch.zeros((N, K, P, Q), device=x.device, dtype=x.dtype).to(memory_format=torch.channels_last) GEMM_M = N * P * Q # 输出元素总数 GEMM_N = K # 输出通道数 GEMM_K = C * R * S # 每个输出元素的乘加操作数 ``` 我们考虑如何在triton中实现这样的一个卷积运算。对于类似Triton这样的GPU编程语言,笔者在这里不予解释的引入`BLOCK_SIZE_M`和`BLOCK_SIZE_N`,他们在GEMM_M和GEMM_N的维度对数据进行了切分,那么我们所有的切分数量为:`triton.cdiv(GEMM_M, META['BLOCK_SIZE_M']) * triton.cdiv(GEMM_N, META['BLOCK_SIZE_N'])` 通过上述切分,我们在核函数程序内计算当前线程块负责计算的部分,其中笔者实现了一种分组策略,目的是优化L2缓存的使用。它将输出矩阵划分为多个组,每个组内的线程块共享相同的输入数据,从而提高缓存效率。: ```python pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(GEMM_M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(GEMM_N, BLOCK_SIZE_N) num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_SIZE_M group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m ``` 随后将计算的结果索引到输出上。`gemm_i`和`gemm_j`分别对应输出矩阵的行和列索引。: ```python gemm_i = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % GEMM_M gemm_j = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % GEMM_N ``` 最关键的步骤在于,我们要把GEMM索引映射回卷积的多维索引。 ```python n = gemm_i // (out_height * out_width) # 批次索引 npq_residual = gemm_i % (out_height * out_width) p = npq_residual // out_width # 输出高度索引 q = npq_residual % out_width # 输出宽度索引 k = gemm_j # 输出通道索引 ``` 这里,我们将一维的`gemm_i`索引解码为卷积输出的三维索引`(n, p, q)`,而`gemm_j`直接对应输出通道`k`。 随后我们便可以按照GEMM乘法的思想,正确计算输入输出并写到特定的索引上。需要注意的是,由于上面的索引粗暴的读取block_size个数据,所以我们在实际的计算过程中还需要正确的设置掩码,把不需要的数据设置为0. Last modification:September 22, 2024 © Allow specification reprint Support Appreciate the author AliPayWeChat Like If you think my article is useful to you, please feel free to appreciate
3 comments
你的文章让我学到了很多知识,非常感谢。 http://www.55baobei.com/Y9VZUgVLkr.html
你的文章让我感受到了不一样的风景,谢谢分享。http://www.fengliangshengwang.com
你的文章让我感受到了不一样的风景,谢谢分享。http://www.fengliangshengwang.com