关于人工智能:旷视MegEngine是如何将3131的大核卷积计算速度提高10倍的

1次阅读

共计 3352 个字符,预计需要花费 9 分钟才能阅读完成。

2021 年提议的 Vision Transformer(VIT)已成为计算机视觉深度学习畛域的一个有前途的钻研主题。随着 VIT 的钻研变得更加深刻,一些钻研人员受到 VIT 的大型接管畛域的启发,将卷积网络也革新成具备更大的接管场来提高效率。依据无效的接管场(ERF)实践,ERF 的大小与内核大小和模型深度的平方根成正比。这意味着通过重叠层来实现大 ERF 的有效性不如减少卷积内核大小。因而,钻研人员提出了蕴含大型卷积内核新的 CNN 构造。该网络能够达到与 VIT 雷同的准确性。“大内核的 cnn 能够胜过小内核的 cnn”这可能是往年来对于 CNN 钻研最大的成绩了。

内核带来了更多的计算和参数

CNN 中最常见的卷积内核大小是 3 ×3、5×5、7×7。个别状况下如果卷积内核的大小超过 9 ×9 就将其确定为“大”。卷积内核的大小减少了 n,参数数量和浮点操作(FLOPS)的数量将大概减少 n 平方。这就是钻研人员在摸索 CNN 新构造时首先不思考大型内核卷积的次要起因。如下图所示,逐通道卷积 (又称深度卷积) 能够将 FLOPs 和参数数量缩小到密集卷积的 1 /(输出通道数量)。这就是为什么大多数钻研人员将大核卷积设计为深度卷积,这样既能够取得大核卷积的益处,又不须要大幅减少参数和 FLOPs。

然而在雷同的 FLOPs 限度下,钻研人员发现,他们的大核深度卷积神经网络比传统的网络运行得慢得多。

MegEngine 团队发现,通过适当的优化,当内核尺寸越来越大时,深度卷积简直不须要任何老本。本文就是对他们研究成果的介绍。

如何寻找大核卷积的优化空间?

为了答复这个问题,咱们须要将 Roofline 模型作为背景。如下图所示,roofline 模型用于形容在计算平台的算力和带宽的限度下,程序所能达到的实践性能上界。

roofline 模型的示意图,有三个重要概念:

  • 算力:每秒所实现的浮点运算次数,单位为 FLOP/ s 或 GFLOP/s
  • 带宽:每秒所实现的内存读取量,单位为 Byte/ s 或 GByte/s
  • 计算密度:又称访存比,是算力与带宽的比值,即每字节读取所实现的浮点运算量,单位为 FLOP/Byte

“roofline”是指 TP 对 IM 的图的形态。设施的实践峰决定了“屋顶”的高度(蓝线)。同时,“屋顶”(红线)的斜率是设施的最大拜访带宽。这两条黑白线的连贯也分隔两种类型的利用,计算联合和内存联合的瓶颈的拆散。

当应用程序的计算密度 I 超过最大 IM 时,其性能只能达到计算设施的实践峰值 TP。应用程序的最大性能 P 显然与计算密度 I 成正比。这种类型的应用程序称为计算绑定。当应用程序的计算密度小于 IM 时,性能是由设施的最大带宽和计算密度确定的。这种类型的应用程序称为内存绑定。与计算绑定计划相同,减少设施带宽或减少内存绑定应用程序的计算密度可能会导致应用程序性能的线性减少。

针对密集卷积的优化曾经达到了瓶颈

到当初为止,很多的钻研曾经摸索对密集卷积的优化技术,包含 Direct, im2col/implicit GEMM, Winograd and FFT。

在本文中应用 NVIDIA 2080TI GPU 作为计算设施在 roofline 模型的背景下进行剖析。该计算设施的 L2 高速缓存带宽为 2.16TB/s,实践峰值性能为 4352 FFMA 外围 1.545 GHz 2 = 2 = 13.447 tflops。咱们假如 CUDA 中每个线程的输入数据都积攒在寄存器中,L1 缓存达到 100%,同时疏忽了写回输入的过程。因为古代计算设施的设计足以在理论卷积计算中同时反对许多耗时的拜访操作,还假如 L2 缓存达到 100%,并达到 L2 缓存的最大带宽。本文应用的卷积输出形态为(n, ic, ih, iw)。内核是(oc, ic, kh, kw),输入是(n, oc, oh, ow)。

im2col/implicit GEMM 是最经典的密集卷积优化办法,在 im2col 转换之后,将卷积转换为一个矩阵乘法问题,其中𝑀=oc,𝑁=𝑛×oh×ow,𝐾=ic×kh×kw,如下图所示。

矩阵乘法在 cuBLAS 等计算库中曾经失去了很好的优化。特地是当矩阵足够大时,性能能够靠近设施的实践峰。咱们在此简要剖析了应用 Roofline 模型的性能。

为了充分利用硬件构造,通常将矩阵乘法的计算划分为块,这样能够使多级存储可能满负荷工作,从而取得最大的内存拜访带宽。

如下图所示,给定 CUDA 中的每个 Thread Block 解决 BM×BN 的输入

  • 内核块大小为 BM×BK
  • 输出块大小为 BK×BN
  • 计算的次数是 BM×BN×BK×2
  • 内存拜访为(BM×BK+BN×BK)×4
  • 计算密度为 BM×BN×2/(BM+BN)×4

依据 Roofline 模型的形容,设施的 IM 计算为 IM=TP/𝐵=13.447/2.16=6.225 FLOPs/byte。为了达到硬件的实践峰值,咱们只须要保障计算密度大于 IM 即可。设 BM=32, BN=32,则计算密度达到 8 FLOPs/byte,显著大于 IM。然而该应用程序属于 Compute Bound 区域,这意味着没有性能优化的空间。咱们能够必定地得出 VGG 时代的旧论断:“随着内核尺寸的增长,工夫也是线性的减少”。

Depthiswise Conv 不是同一状况

随着内核大小的增长,计算的数量随着密度或深度的增长而增长。因为 Cudnn 在 Depthiswise Conv 上的性能趋势的确显示出相似的论断,因而许多人认为先前对于密集的 Conv 的论断在 Depthiswise Conv 上也无效。

这其实是谬误的演绎总结,因为 Depthiswise Conv 的前提条件就核传统的密集卷积不一样。

咱们用 im2col/implicit GEMM 办法剖析 Depthiswise Conv。因为 Depthiswise Conv 会逐通道的计算,因而能够将其视为一组单通道卷积,通道数等于组的大小。在 IM2COL 转换之后,咱们将取得一个批处理的 GEMV,对于下图所示。

放弃与密集块雷同的块大小,每个批的 GEMV 如下图所示。此时的计算密度为 BN×2/(1+𝐵𝑁)×4=BN/(2×BN+2)。给定 BN=1,最大计算密度为 0.25 FLOPs/byte,远小于 IM 6.225。这意味着当初在内存限度区。只管有一些办法能够使 GEMV 更快,但“向量 x 矩阵”的布局注定是内存受限的应用程序。

接下来,咱们将剖析 Direct 办法。例如,让一个 warp 的所有 32 个线程都计算输入 ohow,内核大小为 khkw:

  • 浮点运算次数 =oh×ow×khkw×2 FLOPs
  • 内存拜访数 = (kh×kw+(oh+kh−1)×(ow+kw−1))×4 bytes
  • 内核大小:kh×kw
  • 输出大小:(oh+kh−1)×(ow+kw−1)
  • 计算密度 =(oh×ow×kh×kw×2)/{(kh×kw+(oh+kh−1)×(ow+kw−1))×4}

如果每个线程计算 4 个输入,那么当内核大小为 3 ×3 时,一次 warp 能够计算 4 ×32 个输入。此时计算密度为(4×32×3×3×2)/(3×3+6×34)×4=2.7 FLOPs/byte。最大性能为 2.16*2.7 = 5.84 TFLOPS 远低于 13.447 TFLOPS 的实践峰值。所以为了减少计算密度能够减少每次经纱计算的输入次数。然而因为卷积自身的输入大小以及无限的计算资源(例如每个流多处理器中的寄存器文件),所以并不能有限的减少。

总结一下在 im2col 和 direct 办法中的发现: 深度卷积是一种 Memory Bound 操作。减少内核大小不会大幅扭转内存拜访的数量,计算的工夫应该放弃不变,这就是他与传统的密集卷积最大的区别。然而下面的钻研的好的方向是在适当的优化下,深度卷积的内核大小简直能够随便减少。

MegEngine 的体现

与实践剖析不同的是,在理论利用是发现随着内核大小的减少,cuDNN 的性能会变得很差。

输出形态:(64,384,32,32)

输入形态:(64,384,32,32)

设施:2080 ti

这就是 MegEngine 针对大核深度卷积进行密集优化的起因,如果减少核的大小,须要使计算工夫与实践预测的一样。与 PyTorch 相比,应用 MegEngine 时训练工夫仅为 10%。

https://avoid.overfit.cn/post/abc028ff6f594075a69f177550f56dfc

作者:MagEngine team

正文完
 0