如何实现高速卷积?深度学习库使用了这些「黑魔法」

语言: CN / TW / HK

我的笔记本电脑CPU还可以,在TensorFlow等库的加持下,这台计算机可以在 10-100 毫秒内运行大部分常见CNN模型。2019年,即使是智能手机也能在不到半秒内运行「重量级」CNN模型。而当我自己做了一个简单的卷积层实现,发现这一个层的运行时间竟然超过2秒时,我非常震惊。

大家都知道,现代深度学习库对大部分运算具备高度优化的生产级实现。但是这些库使用了哪些人类不具备的「黑魔法」呢?它们如何将性能提升100倍?当它们「优化」或加速神经网络运算时,它们在做什么?当谈及高性能/高效DNN时,我常常问(或被问及)这些问题。

本文尝试介绍在DNN库中如何实现一个卷积层。卷积不仅是很多DNN模型中最常见和最繁重的操作,它也是导致上述高性能实现的代表性trick:一点点算法智慧 + 大量调参 + 充分利用低级架构。

本文所涉及的很多内容来自这篇开创性论文《Anatomy of High-Performance Matrix Multiplication》,这篇论文形成了线性代数库(如OpenBLAS)所用算法的基础;本文还涵盖 Stefan Hadjis 博士和 Chris Rose的教程。

  • 论文地址:https://www.cs.utexas.edu/~flame/pubs/GotoTOMS_revision.pdf

  • 教程地址:https://cs.stanford.edu/people/shadjis/blas.html

朴素卷积(Naive Convolution)

不成熟的优化是万恶之源。——Donald Knuth

在探讨优化之前,我们先来了解一下基线和瓶颈。这里是一个朴素 numpy/for-loop 卷积:

<code>'''Convolve `input` with `kernel` to generate `output`    </code><code>input.shape = [input_channels, input_height, input_width]    </code><code>kernel.shape = [num_filters, input_channels, kernel_height, kernel_width]    </code><code>output.shape = [num_filters, output_height, output_width]</code><code>'''</code><code>for filter in 0..num_filters    </code><code>    for channel in 0..input_channels        </code><code>        for out_h in 0..output_height            </code><code>            for out_w in 0..output_width                </code><code>                for k_h in 0..kernel_height                    *for* k_w *in* 0..kernel_width                        output[filter, channel, out_h, out_h] +=                             kernel[filter, channel, k_h, k_w] *                             input[channel, out_h + k_h, out_w + k_w]</code>

这个卷积包含6个嵌套的for loop,这里不涉及步幅(stride)、扩张(dilation)等参数。假如这是MobileNet第一层的规模,我们在纯C中运行该层,花费的时间竟然高达22秒!在使用最强悍的编译器优化后(如-O3 或 -Ofast),该卷积层的运行时间降至2.2秒。但是就一个层而言,这个速度仍然太慢了。

那么如果我使用Caffe运行这个层呢?在同一台计算机上使用Caffe运行同一个层所花费的时间仅为18毫秒,实现了100倍的加速!整个网络运行时间才大约100毫秒。

那么「瓶颈」是什么?我们应该从哪儿开始优化呢?

最内的循环进行了两次浮点运算(乘和加)。对于实验所使用的卷积层规模,它执行了8516万次,即该卷积需要1.7亿次浮点运算(170MFLOPs)。我的电脑CPU的峰值性能是每秒800亿FLOPs,也就是说理论上它可以在0.002秒内完成运行。但是很明显我们做不到,同样很明显的是,电脑的原始处理能力是非常充足的。

理论峰值无法实现的原因在于,内存访问同样需要时间:无法快速获取数据则无法快速处理数据。上述高度嵌套的for-loop使得数据访问非常艰难,从而无法充分利用缓存。

贯穿本文的问题是:如何访问正在处理的数据,以及这与数据存储方式有何关联。

先决条件

FLOP/s

性能或者速度的度量指标是吞吐量(throughput),即每秒浮点运算数(FLoating Point Operations per Second,FLOP/s)。使用较多浮点运算的大型运算自然运行速度较慢,因此FLOP/s是对比性能的一种较为可行的指标。

我们还可以用它了解运行性能与CPU峰值性能的差距。我的计算机CPU具备以下属性:

  • 2个物理内核;

  • 每个内核的频率为2.5 GHz,即每秒运行2.5×10^9个CPU循环;

  • 每个循环可处理32 FLOPs(使用AVX & FMA)。

因此,该CPU的峰值性能为:

GFLOP/s。这就是我的CPU的理论峰值性能。类似地,我们可以得出单核CPU的峰值性能为80 GFLOP/s。

存储顺序和行优先

逻辑上我们将矩阵/图像/张量看作是多维度的,但实际上它们存储在线性、一维的计算机内存中。我们必须定义一个惯例,来规定如何将多个维度展开到线性一维存储空间中,反之亦然。

大部分现代深度学习库使用行主序作为存储顺序。这意味着同一行的连续元素被存储在相邻位置。对于多维度而言,行主序通常意味着:在线性扫描内存时第一个维度的变化速度最慢。

那么维度本身的存储顺序呢?通常4维张量(如CNN中的张量)的存储顺序是NCHW、NHWC等。本文假设CNN中的张量使用NCHW存储顺序,即如果HxW 图像的block为N,通道数为C,则具备相同N的所有图像是连续的,同一block内通道数C相同的所有像素是连续的,等等。

Halide语言

本文讨论的优化有时需要使用C语法甚至汇编语言这类底层的低级语言。这会影响代码的可读性,还会使尝试不同优化方法变得困难,因为它需要重写全部代码。Halide是一种嵌入到 C++ 中的语言,它可以帮助抽象概念,旨在帮助用户写出快速的图像处理代码。它可以分离算法(需要计算的东西)和调度策略(如何计算算法以及何时计算),因此使用Halide试验不同优化方法会更加简便。我们可以保持算法不变,试用不用的调度策略。

本文将使用Halide语言展示这些低级概念,但是你需要首先了解函数名称。

从卷积到矩阵相乘

上文讨论的朴素卷积已经够慢了,本节要介绍的实现则更加复杂,它包含步幅、扩张、填充(padding)等参数。我们将继续使用基础卷积作为运行示例,但是最大程度利用计算机的性能需要一些技巧——在多个层次上精心调参,以及充分利用所用计算机架构的具体知识。换言之,解决所有难题是一项艰巨任务。

那么我们可以将它转换成容易解决的问题吗?比如矩阵相乘。

矩阵相乘(又称matmul,或者Generalized Matrix Multiplication,GEMM)是深度学习的核心。全连接层、RNN等中都有它的身影,它还可用于实现卷积。

卷积是滤波器和输入图像块(patch)的点乘。如果我们将滤波器展开为2-D矩阵,将输入块展开为另一个2-D矩阵,则将两个矩阵相乘可以得到同样的数字。与CNN不同,近几十年来矩阵相乘已经得到广泛研究和优化,成为多个科学领域中的重要问题。将图像块展开为矩阵的过程叫做im2col(image to column)。我们将图像重新排列为矩阵的列,每个列对应一个输入块,卷积滤波器就应用于这些输入块上。

下图展示了一个正常的3x3卷积:

下图展示的是该卷积运算被实现为矩阵相乘的形式。右侧的矩阵是im2col的结果,它需要从原始图像中复制像素才能得以构建。左侧的矩阵是卷积的权重,它们已经以这种形式储存在内存中。

注意,矩阵乘积直接得出了卷积输出,无需额外「转换」至原始形式。

出于视觉简洁考虑,此处将每个图像块作为独立的个体进行展示。而在现实中,不同图像块之间通常会有重叠,因而im2col可能导致内存重叠。生成im2col 缓冲(im2col buffer)和过多内存(inflated memory)所花费的时间必须通过GEMM实现的加速来抵消。

使用im2col可以将卷积运算转换为矩阵相乘。现在我们可以使用更加通用和流行的线性代数库(如OpenBLAS、Eigen)对矩阵相乘执行高效计算,而这是基于几十年的优化和微调而实现的。

如果要抵消im2col变换带来的额外工作和内存,我们需要一些加速。接下来我们来看这些库是如何实现此类加速的。本文还介绍了在系统级别进行优化时的一些通用方法。

加速GEMM

朴素

本文剩余部分将假设GEMM按照该公式执行:

我们首先计算基础标准矩阵相乘的时间:

<code>for i in 0..M:    </code><code>    for j in 0..N:        </code><code>        for k in 0..K:            </code><code>            C[i, j] += A[i, k] * B[k, j]</code>

使用Halide语言:

<code>Halide::Buffer C, A, B;</code><code>Halide::Var x, y;</code><code>C(x,y) += A(k, x) *= B(y, k);  // loop bounds, dims, etc. are taken care of automatically</code>

最内的代码行做了2次浮点运算(乘和加),代码执行了M∗N∗K次,因此该GEMM的FLOPs数为2MNK。

下图展示了不同矩阵大小对应的性能:

仅仅达到峰值性能的10%!接下来我们将探索加速计算的方式,但不变的主题仍然是:如果无法快速获取数据,则仅仅快速计算数据并不足够。随着矩阵规模越来越大,内存成为更加严重的问题,而性能也会继续逐渐下降。你看到图中的剧烈下跌了吗?这表示当矩阵太大无法适应缓存时,吞吐量突然下跌。

缓存

RAM是大的存储空间,但速度较慢。CPU缓存的速度要快得多,但其规模较小,因此恰当地使用CPU缓存至关重要。但是并不存在明确的指令:「将该数据加载到缓存」。该过程是由CPU自动管理的。

每一次从主内存中获取数据时,CPU会将该数据及其邻近的数据加载到缓存中,以便利用访问局部性(locality of reference)。

你应该首先注意到我们访问数据的模式。我们按照下图A的形式逐行遍历数据,按照下图B的形式逐列遍历数据。

它们的存储也是行优先的,因此一旦我们找到 A[i, k],则它在该行中的下一个元素A[i, k+1]已经被缓存了。接下来我们来看B中发生了什么:

  • 列的下一个元素并未出现在缓存中,即出现了缓存缺失(cache miss)。这时尽管获取到了数据,CPU也出现了一次停顿。

  • 获取数据后,缓存同时也被 B 中同一行的其他元素填满。我们实际上并不会使用到它们,因此它们很快就会被删除。多次迭代后,当我们需要那些元素时,我们将再次获取它们。我们在用实际上不需要的值污染缓存。

我们需要重新修改loop,以充分利用缓存能力。如果数据被读取,则我们要使用它。这就是我们所做的第一项改变:循环重排序(loop reordering)。

我们将i,j,k 循环重新排序为 i,k,j:

<code>for i in 0..M:    </code><code>    for k in 0..K:        </code><code>        for j in 0..N:</code>

答案仍然是正确的,因为乘/加的顺序对结果没有影响。而遍历顺序则变成了如下状态:

循环重排序这一简单的变化,却带来了相当可观的加速:

平铺(Tiling)

要想进一步改进重排序,我们需要考虑另一个缓存问题。

对于A中的每一行,我们针对B中所有列进行循环。而对于 B 中的每一步,我们会在缓存中加载一些新的列,去除一些旧的列。当到达A的下一行时,我们仍旧重新从第一列开始循环。我们不断在缓存中添加和删除同样的数据,即缓存颠簸(cache thrashing)。

如果所有数据均适应缓存,则颠簸不会出现。如果我们处理的是小型矩阵,则它们会舒适地待在缓存里,不用经历重复的驱逐。庆幸的是,我们可以将矩阵相乘分解为子矩阵。要想计算 C 的r×c平铺,我们仅需要A的r行和B的c列。接下来,我们将 C 分解为6x16的平铺:

<code>C(x, y) += A(k, y) * B(x, k);</code><code>
</code><code>C.update().tile(x, y, xo, yo, xi, yi, 6, 16)</code><code>/*in pseudocode:for xo in 0..N/16:    </code><code>for yo in 0..M/6:        </code><code>    for yi in 6:            </code><code>        for xi in 0..16:                </code><code>            for k in 0..K:                    </code><code>                C(...) = ...*/</code>

我们将x,y 维度分解为外侧的xo,yo和内侧的xi,yi。我们将为该6x16 block优化micro-kernel(即xi,yi),然后在所有block上运行micro-kernel(通过xo,yo进行迭代)。

向量化 & FMA

大部分现代CPU支持SIMD(Single Instruction Multiple Data,单指令流多数据流)。在同一个CPU循环中,SIMD可在多个值上同时执行相同的运算/指令(如加、乘等)。如果我们在4个数据点上同时运行SIMD指令,就会直接实现4倍的加速。

因此,当我们计算处理器的峰值速度时,我们其实有些作弊,把该向量化性能作为峰值性能。对于向量等数据而言,SIMD用处多多,在处理此类数据时,我们必须对每一个向量元素执行同样的指令。但是我们仍然需要设计计算核心,以充分利用它。

计算峰值FLOPs时,我们所使用的第二个技巧是FMA(Fused Multiply-Add)。尽管乘和加是两种独立的浮点运算,但它们非常常见,有些专用硬件单元可以将二者融合为一,作为单个指令来执行。编译器通常会管理FMA的使用。

在英特尔CPU上,我们可以使用SIMD(AVX & SSE)在单个指令中处理多达8个浮点数。编译器优化通常能够独自识别向量化的时机,但是我们需要掌控向量化以确保无误。

<code>C.update().tile(x, y, xo, yo, xi, yi, 6, 16).reorder(xi, yi, k, xo, yo).vectorize(xi, 8)</code><code>/*in pseudocode:for xo in 0..N/16:    </code><code>for yo in 0..M/6:        </code><code>    for k in 0..K:            </code><code>        for yi in 6:                </code><code>            for vectorized xi in 0..16:                    </code><code>                C(...) = ...*/</code><code>
</code>

多线程处理(Threading)

到现在为止,我们仅使用了一个CPU内核。我们拥有多个内核,每个内核可同时执行多个指令。一个程序可被分割为多个线程,每个线程在单独的内核上运行。

<code>C.update().tile(x, y, xo, yo, xi, yi, 6, 16).reorder(xi, yi, k, xo, yo).vectorize(xi, 8).parallel(yo)</code><code>/*in pseudocode:for xo in 0..N/16 in steps of 16:    </code><code>for parallel yo in steps of 6:        </code><code>    for k in 0..K:            </code><code>        for yi in 6:                </code><code>            for vectorized xi in 0..16 in steps of 8:                    </code><code>                C(...) = ...*/</code>

你可能注意到,对于非常小的规模而言,性能反而下降了。这是因为工作负载很小,线程花费较少的时间来处理工作负载,而花费较多的时间同步其他线程。多线程处理存在大量此类问题。

展开(Unrolling)

循环使我们避免重复写同样代码的痛苦,但同时它也引入了一些额外的工作,如检查循环终止、更新循环计数器、指针运算等。如果手动写出重复的循环语句并展开循环,我们就可以减少这一开销。例如,不对1个语句执行8次迭代,而是对4个语句执行2次迭代。

这种看似微不足道的开销实际上是很重要的,最初意识到这一点时我很惊讶。尽管这些循环操作可能「成本低廉」,但它们肯定不是免费的。每次迭代2-3个额外指令的成本会很快累积起来,因为此处的迭代次数是数百万。随着循环开销越来越小,这种优势也在不断减小。

展开是几乎完全被编译器负责的另一种优化方式,除了我们想要更多掌控的micro-kernel。

<code>C.update().tile(x, y, xo, yo, xi, yi, 6, 16).reorder(xi, yi, k, xo, yo).vectorize(xi, 8).unroll(xi).unroll(yi)</code><code>/*in pseudocode:for xo in 0..N/16:    </code><code>for parallel yo:        </code><code>    for k in 0..K:            </code><code>        C(xi:xi+8, yi+0)            </code><code>        C(xi:xi+8, yi+1)            </code><code>        ...            </code><code>        C(xi:xi+8, yi+5)            </code><code>        C(xi+8:xi+16, yi+0)            </code><code>        C(xi+8:xi+16, yi+1)            </code><code>        ...            </code><code>        C(xi+8:xi+16, yi+5)*/</code>

现在我们可以将速度提升到接近60 GFLOP/s。

总结

上述步骤涵盖一些性能加速最常用的变换。它们通常以不同方式组合,从而得到更加复杂的调度策略来计算同样的任务。

下面就是用Halide语言写的一个调度策略:

<code>matrix_mul(x, y) += A(k, y) * B(x, k);    </code><code>out(x, y) = matrix_mul(x, y);</code><code>
</code><code>out.tile(x, y, xi, yi, 24, 32)        </code><code>  .fuse(x, y, xy).parallel(xy)        </code><code>  .split(yi, yi, yii, 4)        </code><code>  .vectorize(xi, 8)        </code><code>  .unroll(xi)        </code><code>  .unroll(yii);</code><code>
</code><code> matrix_mul.compute_at(out, yi)        </code><code>   .vectorize(x, 8).unroll(y);</code><code>
</code><code> matrix_mul.update(0)        </code><code>   .reorder(x, y, k)        </code><code>   .vectorize(x, 8)        </code><code>   .unroll(x)        </code><code>   .unroll(y)        </code><code>   .unroll(k, 2);</code><code>
</code>

它执行了:

  1. 将out分解为32x24的平铺,然后将每个平铺进一步分解为8x24的子平铺。

  2. 使用类似的重排序、向量化和展开,在临时缓冲区(matrix_mul)计算8x24 matmul。

  3. 使用向量化、展开等方法将临时缓冲区matrix_mul 复制回out。

  4. 在全部32x24平铺上并行化这一过程。

最后,我们将速度提升至超过120 GFLOPs,接近峰值性能160 GFLOPs,堪比OpenBLAS等生产级库。使用im2col类似的微调代码和矩阵相乘,同样的卷积可以在大约20毫秒内完成运行。如想深入了解,你可以尝试自行试验不同的调度策略。作为一名工程师,我经常听到关于缓存、性能等的语句,而亲眼看到它们的实际效果是非常有趣和值得的。

注意:im2col+gemm方法是大部分深度学习库中流行的通用方法。但是,对于特定的常用规模、不同的架构(GPU)和不同的运算参数(如扩张、分组等),专门化(specialization)是关键。这些库可能也有更专门化的实现,这些实现利用类似的trick或具体的假设。经过不断试错的高度迭代过程之后,我们构建了这些micro-kernel。对于什么方法效果好/不好、必须基于结果考虑注释,程序员通常只有直观的想法。听起来和深度学习研究差不多,不是吗?

原文链接:

https://sahnimanas.github.io/post/anatomy-of-a-high-performance-convolution/

分享到:
其他文章