为什么是大 kernel 卷积？
Transformer 目前在 CV 畛域愈发炽热，这份炽热促使着优良学者们思考一个更深层次的问题。局部学者认为 Transformer 之所以 work 更加实质的起因在于其大的感触野 （论文中转）。依据无效感触野（ERF）实践，ERF 大小与 kernel 大小成正比关系，与模型深度的平方根也成正比关系。所以通过重叠层数实现大感触野必然不如减少卷积 kernel 大小更高效。因而有学者提出超大 kernel 卷积的网络结构，并证实在指标检测和语义宰割等工作上超过 Swin Transformer 而且远超传统小卷积模型。
什么是大 kernel，什么是 depthwise 卷积？
CNN 中最常见的卷积 kernel 大小有 2×2, 3×3, 5×5, 7×7 等，在本文中咱们将卷积 kernel 大小超过 9×9 的视作大 kernel，同时以下所有数据都是近似数据。咱们不难看出随着卷积 kernel 大小的减少，卷积的参数量和计算量都呈平方增长，这往往也是大家不喜爱用大 kernel 卷积的其中一个起因。为了取得大 kernel 卷积带来的收益的同时升高其计算量和参数量，咱们个别将大 kernel 卷积设计成 depthwise 卷积。如下图所示，depthwise 卷积通过逐通道（channel）做卷积，能够将计算量和参数量升高到 Dense 卷积的 input channel 分之一。
大 kernel depthwise 卷积为什么值得优化？
Roofline Model
为了解释分明为什么大 kernel 值得优化这个问题，咱们须要借助 Roofline 模型的帮忙。如下图所示，Roofline 尝试解释一件非常简单的事件，即利用在特定计算设施下能达到多快的计算速度。
- 实践峰值 TP：形容了计算设施的性能下限，指的是一个计算设施每秒钟最多所能实现的浮点运算数，单位是
FLOPS。 - * 最大带宽 B8：形容计算设施的带宽下限，指的是一个计算设施每秒最多所能实现的内存交换量，单位是
Byte/s。 - 最大计算密度 IM：形容计算设施单位内存替换最多用来进行多少次运算，单位是
FLOPs/Byte。
“Roofline” 指的是由计算设施实践算力峰值和最大访存带宽这两个参数所决定的“屋顶”状态。其中设施实践峰值决定“屋顶”的高度（蓝色线段），设施最大访存带宽决定了“屋檐”的斜率（红色线段）。Roofline 模型划分进去两个瓶颈区域，别离为 Compute Bound 和 Memory Bound。
当利用的计算密度 I 超过最大计算密度 IM 时，此时无论利用的计算密度多大，它的性能最高只能达到计算设施的实践峰值 TP。此时利用的性能 P 被设施实践峰值限度无奈和计算密度 I 成正比，所以叫做 Compute Bound。当利用的计算密度 I 小于最大计算密度 IM 时，此时性能 P 将由设施最大带宽和利用计算密度决定。不难看出对于处在 Memory Bound 区间的利用，减少设施带宽和减少计算密度能够使利用性能达到线性增长的目标。
走出对 depthwise 卷积速度的 “ 思维误区 ”
为什么不是大 kernel Dense 卷积
现如今针对 Dense 卷积咱们曾经有了包含 Direct、im2col/implicit GEMM、Winograd 和 FFT 等多种优化伎俩，能够说曾经足够成熟了。可是如果咱们抛开模型参数量，仅仅从运行效率的角度思考一个问题，为什么咱们不必大 kernel Dense 卷积而抉择大 kernel depthwise 卷积呢？
为了探寻这个问题的答案，咱们联合 Roofline 模型具体分析。本文选取 2080Ti 显卡为计算设施，它的实测 L2 cache 带宽为 2.16TB/s，实践峰值性能为 4352 FFMA Cores 1.545 GHZ 2 = 13.447 TFLOPS。咱们假如 CUDA 中每个 thread 负责计算的 output 数据都放在寄存器中累加，咱们假如 L1 cache 100% 命中，疏忽写回 output 的过程。因为古代计算设施的设计足够正当，理论卷积计算中足以对消很多耗时较长的访存操作，同时为了简化剖析复杂度，在这里咱们假如 L2 cache 100% 命中，应用 L2 cache 的最大带宽作为剖析参数。本文应用的卷积输出 shape 是（n, ic, ih, iw），kernel 是（oc, ic, kh, kw），output 是（n, oc, oh, ow）。
对 Dense 卷积而言，一种通用优化计算伎俩就是 im2col/implicit GEMM。因为其太经典了咱们在这里不再赘述 im2col 的过程，感兴趣的能够翻阅咱们之前写的文章《MegEngine TensorCore 卷积算子实现原理》。在通过了 im2col 变换之后，咱们就胜利的将卷积转换成了矩阵乘的模式。其中矩阵乘的 M = oc, N = n*oh*ow, K = ic*kh*kw，具体如下图所示。
对于矩阵乘特地是大规模矩阵乘，cuBlas 等计算库曾经优化的足够好了，基本上能够靠近设施实践峰值，这里咱们联合 Roofline 简略剖析一下性能。为了充沛适应硬件体系结构特色，充分利用多级存储增大访存带宽，咱们须要对矩阵乘进行分块计算。如下图所示，如果 cuda 中每个 Thread Block 解决 BMxBN 的 output，此时 kernel 分块大小为 BMxBK，input 分块大小为 BKxBN。则计算量为 BM*BN*BK*2，访存量为 (BM*BK + BN*BK)*4。计算密度为 $\frac{BM*BN*2}{(BM+BN)*4}$
。依照 Roofline 模型的形容，计算设施的 $IM = \frac{TP}{B} = \frac{13.447}{2.16} = 6.225$
FLOPs/Byte，若要达到设施实践峰值咱们只有保障计算密度大于 IM 即可。如果咱们依照 BM=32, BN=32 来算的话，则此时的计算密度将达到 8 FLOPs/Byte，显然是大于 IM 的。此时如果疏忽 TP 的限度如果打满设施最大带宽，最大可能达到的性能 P = 8*2.16 = 17.28 TFLOPS。联合 Roofline 模型不难看出此时处于 Compute Bound 区域。因为 Compute Bound 区域的计算速度曾经靠近实践峰值，曾经不能减少了。如果咱们采纳大 kernel 的话， 随着 kernel size 的减少计算量会呈平方增长，所以相应的运行工夫也会随之增长，这显然是不可承受的 。
depthwise 卷积速度的“骗局”
对 Dense 卷积剖析让咱们失去了一个论断即“随着 kernel 的增大，卷积工夫呈平方增长”。很多人想当然的将这个论断平移到了 depthwise 卷积上，这其实是一种思维误区。
让咱们同样尝试用 im2col/implicit GEMM 的办法剖析 depthwise 卷积。因为 depthwise 是逐 channel 做卷积的，所以能够看做 channel 数量的单通道卷积。在通过 im2col 变换之后咱们将取得一个 Batched GEMV，每个 batch 的 GEMV 如下图所示。
如果咱们放弃和 Dense 卷积一样的分块策略的话，每个 batch 的 GEMV 如下图所示。相应的此时的计算密度为 $\frac{BN*2}{(1+BN)*4} = \frac{BN}{2*BN+2}$。先不说这是一个 Batched GEMV，独自看一个 GEMV 也不难发现此时的计算密度是很差的，BN = 1 时最高大略能达到 0.25 FLOPs/Byte，相应的最大达到的性能 P = 0.25*2.16 = 0.54 TFLOPS。当然了理论利用中 GEMV 还有其余计算形式，咱们的分析方法就不肯定精确了。但此处想表白的意思是 Batched GEMV 比 GEMM 更难优化。如果 kernel 为 3×3，此时 M=1, K=9, N 受限于 oh 和 ow 也不会很大，此时的 GEMV 性能必定远达不到峰值，并且 GEMV 也不能利用 TensorCore 减速。
如果咱们尝试应用 Direct 的形式解决 depthwise 卷积的话会不会好一点呢？例如咱们让 cuda 中每个 warp 32 个线程负责计算 ohxow 的输入，kernel size 为 khxkw，此时：
- 计算量 = oh*ow*kh*kw*2 FLOPs
-
访存量 = (kh*kw + (oh+kh-1)*(ow+kw-1)) * 4 Bytes，别离为
- kernel: kh*kw
- input: (oh+kh-1)*(ow+kw-1)
- 计算密度为 $\frac{oh*ow*kh*kw*2}{(kh*kw+(oh+kh-1)*(ow+kw-1))*4}$
咱们以一个更具体的例子剖析，如果咱们让每个 thread 负责计算 4 个 output 的话，则一个 warp 负责计算 4×32 的 output，以 kernel（3, 3) 为例。则计算密度为 $\frac{4*32*3*3*2}{(3*3+6*34)*4} = 2.7 $ FLOPs/Byte，最大可达到的性能为 2.16*2.7 = 5.84 TFLOPS，相比于实践峰值 13.447 TFLOPS 仍有很大差距。尽管减少 output 能持续减少计算密度，然而受限于卷积自身的输入大小和每个 SM 中无限的 register file 等计算资源，每个 warp 计算的 output 并不能有限减少。这也是 depthwise 卷积须要更加认真的优化，否则一不小心性能就会很差的其中一个起因。
综合 im2col 和 Direct 两个方面的剖析论断，咱们意识到和 Dense 卷积不同的是 depthwise 卷积很多时候是一个 Memory Bound 的操作。而联合 Roofline 模型对 Memory Bound 瓶颈的剖析和倡议，此时减少计算密度和减少带宽都能够减少性能。在固定设备的状况下咱们无奈减少带宽了，所以看起来减少计算密度是一个可行的计划。通过观察计算密度公式咱们不难发现，减少 depthwise 卷积的 kernel size 就是一个减少其计算密度的无效计划，例如放弃每个 warp 4×32 的输入配置下 kernel size 31×31 的 depthwise 卷积计算密度将达到 $\frac{4*32*31*31*2}{(31*31+34*62)*4} = 20$ FLOPs/Byte，不难看出此时曾经变成了 Compute Bound 的操作。
综上所述，减少卷积 kernel size 会使得计算量减少。同时因为 Dense 卷积处于 Compute Bound 区域，所以其运行速度受限于设施实践峰值无奈晋升，因而针对 Dense 卷积咱们不难演绎出 “随着 kernel 的增大，卷积工夫呈平方增长” 的法则。然而 depthwise 卷积是一种 Memory Bound 的操作，而随着 kernel size 的减少其计算密度也会增大，所以其运行性能也会随之增大。此时的卷积的运行工夫并不会显著增长，所以它并不实用 “随着 kernel 的增大，卷积工夫呈平方增长” 这个论断。这也是咱们认为大 kernel depthwise 还有较大的优化后劲，其运行工夫并不会显著差于小 kernel depthwise 卷积的根据。
现有优化办法为什么不行？
上一节咱们曾经解释了为什么 im2col/implicit GEMM 不适宜 depthwise 卷积，direct 也须要付出很大精力能力写好。另外，提到大 kernel 则不能不提 FFT 算法，但 FFT 在计算 depthwise 卷积的时候只能逐通道计算，性能不如预期。并且 FFT 有其缺点例如精度问题，对半精度计算并不敌对，也不能被量化。咱们在 2080Ti 上应用 input 和 output 形态都是 (n, c, h, w) = (64, 384, 32, 32) 的用例对 cudnn 做了一次测速，咱们遍历所有的 cudnn 算子（内含 FFT）并抉择最快的那个算子进行测试。后果如下：
在大 kernel size 下 cudnn 的体现很差，次要起因是 cudnn 没有针对性优化。咱们留神到很多时候 cudnn 调用到了外部的 implicit_gemm 实现，这不利于施展设施的计算性能。因为对于 depthwise 卷积而言，im2col 之后将会是一个 batch = channel，M = 1，N=nhw，K = kh*kw 的 batched GEMV，这种状况也很难打满设施峰值。
MegEngine 的优化成果和简略剖析
鉴于以上剖析，大 kernel depthwise 卷积有很大的优化后劲，所以 MegEngine 紧跟学界动静对大 kernel depthwise 卷积进行了深度优化。如上图所示，通过咱们的优化后，随着 kernel size 的减少，算子性能根本出现线性增长的趋势，局部状况下算子能够迫近硬件的单精度浮点实践峰值。
如下图所示，优化后的大 kernel depthwise 卷积比 PyTorch 快 10.x 倍，代码附在文末，感兴趣的同学欢送来体验一把。而且咱们不难发现，随着 kernel size 的减少模型训练工夫并没有显著减少。起因就在于 kernel size 不够大的时候算子处于 Memory Bound 状态，远没有达到实践峰值，此时减少计算密度反而不会对算子运行工夫造成很大影响。
想晓得 MegEngine 是如何将 31*31 的 DWconv 优化快了 10 余倍？还有 ConvNext，RepLKNet 为何不谋而合将 kernel size 增大，更大的 kernel size 到底给模型带来了什么？来 MegEngine Meetup 一起聊聊吧。
3.19 日 Meetup 预报
北京工夫本周六（3.19）上午 10:00，MegEngine Meetup 围绕“Large Kernel Makes CNN Great Again”主题，将为大家带来精彩线上分享。
流动信息：周六直播预报 | 突破思维惯性，旷视 MegEngine 通知你为什么要思考大 kernel size
直播间地址:https://live.bilibili.com/224…
附：测试代码
MegEngine 测试代码
import time
import megengine.module as M
import megengine.autodiff as ad
import megengine
import numpy as np
megengine.functional.debug_param.set_execution_strategy("PROFILE")
def benchmark_lknet(ksize, batch=64, dim=384, res=32, depth=24):
m = M.Sequential(*[M.Conv2d(dim, dim, ksize, padding=ksize//2, groups=dim, bias=False) for _ in range(depth)]
)
x = megengine.tensor(np.ones([batch, dim, res, res]))
gm = ad.GradManager().attach(m.parameters())
for i in range(20):
t = time.perf_counter()
with gm:
y = m(x)
gm.backward(y.mean())
megengine._full_sync()
t = time.perf_counter() - t
if i > 9 and i % 10 == 0:
print(t)
return t
if __name__ == "__main__":
args = dict()
for k in (3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31):
t = benchmark_lknet(k, **args)
print("kernel size", k, "iter time", t * 1000, "ms")PyTorch 测试代码
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import time
cudnn.benchmark = True
def benchmark_lknet(ksize, batch=64, dim=384, res=32, depth=24):
m = nn.Sequential(*[nn.Conv2d(dim, dim, ksize, padding=ksize//2, groups=dim, bias=False) for _ in range(depth)]
).cuda()
x = torch.rand(batch, dim, res, res).cuda()
for i in range(20):
t = time.perf_counter()
y = m(x)
y.mean().backward()
torch.cuda.synchronize()
t = time.perf_counter() - t
if i > 9 and i % 10 == 0:
print(t)
return t
if __name__ == "__main__":
args = dict()
for k in (3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31):
t = benchmark_lknet(k, **args)
print("kernel size", k, "iter time", t * 1000, "ms")GitHub：MegEngine 天元
官网：MegEngine- 深度学习，简略开发