paper:Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs
official implementation:https://github.com/DingXiaoH/RepLKNet-pytorch
背景
卷积神经网络(CNN)曾经是现代计算机视觉系统中的常见选择。然而,最近视觉变换器(ViT)在许多视觉任务中表现出了领先的性能,这对CNN构成了巨大挑战。研究人员普遍认为,ViT的强大性能部分归功于其多头自注意力机制(MHSA),该机制在建构大感受野方面具有重要作用。
出发点
本文的出发点是探讨在现代CNN中引入大卷积核的设计,试图通过使用大卷积核来缩小CNN与ViT之间的性能差距。
创新点
本文解决了如何设计有效且高性能的大卷积核CNN架构的问题,特别是在与ViTs性能差距的缩小上。具体包括:
- 大卷积核设计:提出了使用大深度卷积核(最大为31×31)来构建大感受野的CNN架构,这与传统的使用小卷积核堆叠的方法形成鲜明对比。
- 重新参数化技术:通过小卷积核的重新参数化来弥补优化问题,提升了大卷积核CNN的实际效率。
- RepLKNet架构:提出了一种纯CNN架构,采用了大卷积核来构建大的感受野(ERFs),并展示了与ViTs相比具有可比性或更优的结果。
- 五个设计准则:总结了有效使用大卷积核的五个经验性准则,包括应用重参数化的大核深度卷积、identity shortcut connection的重要性、小卷积核的重参数化帮助优化问题、大卷积核对下游任务的提升远超过ImageNet分类,以及大卷积核即使在小特征图上也是有用的。
- 性能提升:RepLKNet在ImageNet分类和一些经典的下游任务上取得了与Swin Transformer相当或更好的结果,同时具有更低的延迟。
- 可扩展性:RepLKNet在大型数据集和大型模型上表现出良好的可扩展性,获得了87.8%的ImageNet top-1准确率和56.0%的ADE20K mIoU,这在类似大小的SOTA模型中非常有竞争力。
- 形状偏好:研究表明,与小卷积核CNN相比,大卷积核CNN具有更大的有效感受野和更高的形状偏好而非纹理偏好, 这与人类的认知部分一致。
方法介绍
Guidelines of Applying Large Convolutions
作者首先总结了有效使用大核卷积的5个指导方针。
Guideline 1:在实际应用中,大核深度卷积是有效的。大家普遍认为大核卷积在计算上是昂贵的,因为核的大小二次方的增加了参数量和FLOPs。通过应用深度卷积(depthwise convolution, DW)可以极大地克服这个缺点。例如在本文提出的RepLKNet上将不同stage的卷积核从[3, 3, 3, 3]增大到[31, 29, 27, 13]只将FLOPs和参数量分别增加了18.6%和10.4%,这是可以接受的。剩下的1x1卷积实际上占据了大部分的复杂度。
有人可能会担心DW卷积在像GPU这样的现代并行计算设备上非常低效。对于传统的3x3 DW卷积确实如此,因为DW操作使得计算 vs. 内存访问成本的比率较低,这对现代计算结构不友好。但是作者发现,当kernel size变大时计算密度会增加:例如在DW 11x11 kernel中,每次我们从特征图中加载一个值,它最多可以进行121次乘法,而在3x3 kernel中这个数字只有9。所以根据roofline model,当卷积核变大时,实际延迟的增加不会像FLOPs增加的那么多。
Remark 1. 不幸的是现有的深度学习框架(例如Pytorch)对深度卷积的支持很差,如表1所示。因此作者尝试了几种方法来优化CUDA kernel,其中基于FFT的方法在实现大核卷积方面似乎是合理的选择,但实践中作者发现block-wise (inverse) implicit gemm算法是更好的选择,如表1所示,和pytorch baseline相比,本文的实现效率要高得多。
Guideline 2:对于大核卷积的网络,identity shortcut至关重要。为了验证这一点,使用MobileNet v2进行基准测试,因为它大量使用DW卷积,并且有两个版本(有shortcut的和没有shortcut的)。作者用13x13 DW卷积替换所有的3x3 DW卷积,并在ImageNet上训练了100个epoch,结果如表2所示,大核卷积精度提升了0.77%,但如果没有shortcut,精度降低到了53.98%。
Remark 2. 这条准则也适用于ViTs。最近一项研究表明,如果没有identity shortcut,注意力会随着深度呈双指数下降,导致过渡平滑问题。虽然大核CNNs可能以不同的机制退化,但我们也观察到,如果没有identity shortcut,网络很难捕获到局部细节。
Guideline 3:用小核重参数化有助于弥补优化问题。如表3所示,直接将卷积核从9x9增大到13x13精度会下降,但额外增加一个3x3并行分支并在推理时进行重参数化解决了这个问题,如图2所示。
Remark 3. 众所周知ViTs存在优化问题特别是在小数据集上,一个常见的解决方法是引入卷积先验,例如在每个self-attention block中添加一个3x3 DW conv,这和我们的方法类似。这些策略向网络中引入了平移不变性和局部先验,使得网络在小数据集上更容易优化。与ViT类似,我们还发现,当数据量增加到7300万张图片时,重参数化可以省略。
Guideline 4:相比于ImageNet分类,大核卷积更能提升下游任务的性能。如表3所示,卷积核从3x3增大到9x9,ImageNet上提升了1.33%,而Cityscapes上mIoU提升了3.99%。表5也显示出类似的趋势,当卷积核从[3, 3, 3, 3]增大到[31, 29, 28, 13]时,ImageNet上的精度只提升了0.97%,但ADE20K上mIoU提升了3.12%。
Remark 4. 是什么导致了这种现象?首先大卷积核显著增加了有效感受野(ERFs),许多工作已经证明,上下文信息即大ERFs,对目标检测和语义分割等下游任务是至关重要的。其次,作者任务大核卷积让网络有了更多了shape bias,简单地说,ImageNet可以通过纹理或形状进行分类。但人识别物体更多的是基于形状而不是纹理,因此具有较强shape bias的模型可能会更好的转移到下游任务。
Guideline 5:即使在小的特征图(例如7x7)上大核卷积(例如13x13)也是有用的。为了验证这一点,作者将MobileNet v2最后一个stage的卷积扩大到7x7或13x13,这样卷积核就与特征图(默认7x7)相当大小甚至更大。结果如表4所示,即使最后一个stage卷积已经具有很大的感受野了,进一步增大卷积核大小还能提升性能,尤其是在下游任务中。
Remark 5. 注意当kernel size变大时,CNN的平移不变性并不严格成立。如图3所示,相邻位置的两个输出只共享一小部分权重,即它们是通过不同的映射转换得到的。这也符合ViTs的特性,放松对称性以获得更多的容量。
RepLKNet:A Large-Kernel Architecture
基于上述准则,作者提出了RepLKNet,一个大核纯CNN结构设计的网络。整体结构如图4所示。
Stem部分首先是一个3x3-s2的普通卷积进行降采样,然后是一个3x3 DW用来捕捉low-level patterns,接着是一个1x1卷积,然后是另一个3x3-s2 DW进行降采样。
Stage 1-4每个都包括若干RepLK Block,其中使用了shortcut(Guideline 2)和大核DW卷积(Guideline 1)。遵循常见做法,在DW conv前后都使用了1x1 conv。注意每个大核DW卷积都使用了一个5x5 kernel用来重参数化(Guideline 3),在图4中没有显示。受Transformers和MLPs中广泛使用的Feed-Forward Network(FFN)的启发,作者也使用了一个类似的CNN-style block,包含shortcut、BN、两个1x1卷积和GELU,称之为ConvFFN Block。和FFN中使用的LN相比,BN的优点是可以融入前面的卷积中从而进行高效的推理。在网络中,每个RepLK Blcok前都有一个ConvFFN。
Transition Blocks位于stage之间,首先通过一个1x1卷积进行升维,然后通过3x3 DW卷积降采样。
总的来说,每个stage有三个结构参数,RepLK Block的数量 \(B\),通道维度 \(C\),kernel大小 \(K\)。所以一个RepLKNet结构通过 \([B_1,B_2,B_3,B_4],[C_1,C_2,C_3,C_4],[K_1,K_2,K_3,K_4]\) 来定义。
作者通过固定 \(\mathbf{B}=[2,2,18,2],\mathbf{C}=[128,256,512,1024]\) 并改变 \(\mathbf{K}\) 的大小来评估不同大小kernel的表现。作者设置kernel sizes分别为 \([13,13,13,13],[25,25,25,13],[31,29,27,13]\) 得到三个模型variants并表示为RepLKNet-13/25/31。此外作者还构建了两个小核baseline卷积核分别为3和7(RepLKNet-3/7)。
表5是不同kernel size的RepLKNet在ImageNet和ADE20K上的性能,卷积核从3增大到13提升了精度,进一步增大并不会带来进一步的提升了。但是在ADE20K上,将卷积核从[13, 13, 13, 13]增大到[31, 29, 27, 13]将mIoU提升了0.82,表明了大核在下游任务中的重要性。
接下来作者用更强的训练配置训练RepLKNet-31来和其它SOTA模型比较,上述的模型称为RepLKNet-31B(Base),并构建了一个更宽的模型 \(\mathbf{C}=[192,384,768,1536]\) 的RepLKNet-31L(Large)和 \(\mathbf{C}=[256,512,1024,2048]\) 的RepLKNet-XL(DW conv的输出通道数是输入的1.5倍)。
实验结果
由于RepLKNet的整体结构和Swin Transformer类似,作者在ImageNet上比较了两者的精度,如表6所示。可以看到尽管大核卷积对ImageNet分类的帮助没有那么大,但RepLKNet显示出了良好的精度/效率的平衡。例如RepLKNet-31B取得了84.8%的精度比Swin-B搞了0.3%,同时快了43%。RepLKNet-XL即使FLOPs大于Swin-L,但推理速度更快。
作者使用在ImageNet上预训练的模型作为backbone,分割模型采用UperNet,然后比较了在Cityscapes和ADE20K上的表现,如表7、8所示。在Cityscapes上,在ImageNet-1K上预训练的RepLKNet-31B的表现远好于Swin-B,甚至比在ImageNet-22K上预训练的Swin-L还要好。而在ADE20K上,同在ImageNet-22K上预训练的RepLKNet-31B和31L也要好于Swin-B和Swin-L。
表9展示了在COCO数据集上目标检测的表现,可以看到RepLKNet相比于Swin Transformer精度更高同时参数量和FLOPs都更低。