原文题目:《SegNeXt: Rethinking Convolutional Attention Design for Semantic Segmentation》
原文引用:Guo M H, Lu C Z, Hou Q, et al. Segnext: Rethinking convolutional attention design for semantic segmentation[J]. Advances in Neural Information Processing Systems, 2022, 35: 1140-1156.
原文链接:https://proceedings.neurips.cc/paper_files/paper/2022/file/08050f40fff41616ccfc3080e60a301a-Paper-Conference.pdfhttps://proceedings.neurips.cc/paper_files/paper/2022/file/08050f40fff41616ccfc3080e60a301a-Paper-Conference.pdf
我们提出了SegNeXt,这是一种简单的卷积网络架构,用于语义分割。最近基于transformer的模型在语义分割领域占据主导地位,这是由于自注意力在编码空间信息方面的高效性。在本文中,我们展示了卷积注意力是一种比transformer中的自注意力机制更高效和有效地编码上下文信息的方式。通过重新审视成功的分割模型所拥有的特性,我们发现了几个关键组件,这些组件导致了分割模型性能的提升。这激励我们设计了一种新颖的卷积注意力网络,它使用廉价的卷积操作。没有花哨的东西,我们的SegNeXt在流行的基准测试中显著提高了以前最先进方法的性能,包括ADE20K、Cityscapes、COCO-Stuff、Pascal VOC、Pascal Context和iSAID。值得注意的是,SegNeXt在仅使用其1/10参数的情况下超过了EfficientNet-L2 w/ NAS-FPN,在Pascal VOC 2012测试排行榜上达到了90.6%的mIoU。平均而言,SegNeXt在ADE20K数据集上与最先进方法相比,mIoU提高了约2.0%,计算量相同或更少。
作为计算机视觉中最基础的研究课题之一,语义分割旨在为每个像素分配一个语义类别,在过去的十年中引起了广泛关注。从早期的基于CNN的模型,如FCN和DeepLab系列,到最近的基于transformer的方法,如SETR和SegFormer,语义分割模型在网络架构方面经历了重大革命。
通过重新审视以前成功的语义分割作品,我们总结了一些不同模型具有的关键特性,如表1所示。基于以上观察,我们认为一个成功的语义分割模型应具备以下特点:
(i) 强大的主干网络作为编码器。与以前的基于CNN的模型相比,基于transformer的模型的性能改进主要来自更强大的主干网络。
(ii) 多尺度信息交互。不同于主要识别单个对象的图像分类任务,语义分割是一项密集预测任务,因此需要在单个图像中处理不同大小的对象。
(iii) 空间注意力。空间注意力允许模型通过对语义区域内的区域进行优先级排序来进行分割。
(iv) 低计算复杂度。当处理来自遥感和城市场景的高分辨率图像时,这一点尤为重要。
考虑到上述分析,在本文中,我们重新思考了卷积注意力的设计,并提出了一种高效而有效的用于语义分割的架构。与以前的基于transformer的模型在解码器中使用卷积作为特征优化器不同,我们的方法颠倒了transformer-卷积编码器-解码器的架构。具体而言,对于我们编码器中的每个块,我们改进了传统卷积块的设计,并利用多尺度卷积特征通过简单的逐元素乘法来唤起空间注意力,这是在[25]之后的一种方法。我们发现,这种简单的构建空间注意力的方式比标准卷积和自注意力在空间信息编码方面更有效率。对于解码器,我们从不同阶段收集多级特征,并使用Hamburger [22]进一步提取全局上下文。在这种设置下,我们的方法可以从局部到全局获取多尺度的上下文,实现在空间和通道维度上的适应性,并从低级别到高级别聚合信息。
我们的网络被称为SegNeXt,除了解码器部分外,它主要由卷积操作组成,解码器部分包含一个基于分解的Hamburger模块[22](Ham)用于全局信息提取。这使得我们的SegNeXt比之前严重依赖于transformer的分割方法更加高效。如图1所示,SegNeXt在性能上显著优于最近的基于transformer的方法。特别是,在处理来自Cityscapes数据集的高分辨率城市场景时,我们的SegNeXt-S在仅使用约1/6(124.6G vs. 717.1G)的计算成本和1/2(13.9M vs. 27.6M)的参数时,表现优于SegFormer-B2(81.3% vs. 81.0%)。
我们的贡献可以总结如下:
表1:我们从成功的语义分割方法中观察到的对模型性能提升有益的属性。这里,n表示像素或令牌的数量。强大的编码器表示强大的主干网络,如ViT [16]和VAN [25]。
图1:在Cityscapes(左图)和ADE20K(右图)验证集上的性能-计算曲线。FLOPs是使用2,048×1,024的输入大小计算的Cityscapes和512×512的ADE20K。圆的大小表示参数的数量。更大的圆表示更多的参数。我们可以看到,我们的SegNeXt在分割性能和计算复杂度之间实现了最佳的权衡。
语义分割是一项基本的计算机视觉任务。自从FCN [60]被提出以来,卷积神经网络(CNNs)[1, 71, 98, 106, 20, 99, 79, 21, 51]取得了巨大的成功,并成为语义分割的流行架构。最近,基于transformer的方法[108, 90, 100, 73, 70, 50, 10, 9]展现了巨大的潜力,并超过了基于CNN的方法。
在深度学习时代,分割模型的架构可以大致分为两部分:编码器和解码器。对于编码器,研究人员通常采用流行的分类网络(如ResNet [28]、ResNeXt [91]和DenseNet [33])而不是定制的架构。然而,语义分割是一种密集预测任务,与图像分类不同。在分类任务中的改进可能在具有挑战性的分割任务中并不明显 [29]。因此,出现了一些定制的编码器,包括Res2Net [21]、HRNet [79]、SETR [108]、SegFormer [90]、HRFormer [100]、MPViT [44]、DPT [70]等。对于解码器,它通常与编码器合作以实现更好的结果。针对不同的目标,有不同类型的解码器,包括实现多尺度感受野 [106, 6, 88]、收集多尺度语义 [71, 90, 7]、扩大感受野 [4, 4, 69]、加强边缘特征 [107, 2, 15, 48, 102]和捕捉全局上下文 [20, 35, 101, 46, 24, 27, 103]。
在本文中,我们总结了那些成功的用于语义分割的模型的特点,并提出了一种基于CNN的模型,名为SegNeXt。与我们的论文最相关的工作是[69],它将一个k×k的卷积分解为一个k×1和一个1×k的卷积。尽管这项工作显示了在语义分割中大型卷积核的重要性,但它忽略了多尺度感受野的重要性,并且没有考虑如何利用通过大型卷积核提取的这些多尺度特征以注意力的形式进行分割。
设计多尺度网络是计算机视觉中的一个热门方向。对于分割模型来说,多尺度块出现在编码器[79, 21, 75]和解码器[106, 98, 5]部分。GoogleNet [75]是与我们的方法最相关的多尺度架构之一,它使用多分支结构实现多尺度特征提取。另一个与我们的方法相关的工作是HRNet [79]。在更深层次上,HRNet也保持高分辨率特征,并与低分辨率特征聚合,以实现多尺度特征提取。
与之前的方法不同,SegNeXt除了在编码器中捕捉多尺度特征外,还引入了一种高效的注意力机制,并采用了更便宜且更大的卷积核。这使得我们的模型能够实现比前述分割方法更高的性能。
注意力机制是一种自适应选择过程,旨在使网络专注于重要的部分。一般来说,在语义分割中,它可以分为两类 [26],包括通道注意力和空间注意力。不同类型的注意力起着不同的作用。例如,空间注意力主要关注重要的空间区域 [16, 13, 64, 58, 23]。而使用通道注意力的目标是使网络有选择地关注那些重要的对象,在先前的工作中已经证明了它的重要性 [31, 8, 80]。谈到最近流行的视觉Transformer [16, 58, 94, 81, 82, 57, 90, 34, 56, 100, 93],它们通常忽视了在通道维度上的适应性。
视觉注意力网络(VAN)[25]是与SegNeXt最相关的工作,它也提出利用大型卷积核注意力(LKA)机制构建通道和空间注意力。尽管VAN在图像分类方面取得了很好的性能,但它在网络设计过程中忽视了多尺度特征聚合的作用,而这对于分割类任务来说是至关重要的。
图2:所提出的MSCA和MSCAN的示意图。这里,d,k1×k2表示使用k1×k2的核大小进行深度卷积(d)。我们使用卷积提取多尺度特征,然后将它们用作MSCA输入的注意力权重进行重新加权。
在这一部分,我们详细描述了所提出的SegNeXt的架构。基本上,我们采用了一个编码器-解码器的架构,这是大多数先前工作所采用的,简单易懂。
我们采用了金字塔结构作为我们的编码器,这是根据大多数先前的工作[90, 4, 20]。对于我们编码器中的构建模块,我们采用了类似于ViT [16, 90]的结构,但不同之处在于我们没有使用自注意机制,而是设计了一个新颖的多尺度卷积注意力(MSCA)模块。如图2(a)所示,MSCA包含三个部分:一个深度卷积用于聚合局部信息,多分支深度条带卷积用于捕捉多尺度上下文,以及一个1×1卷积用于建模不同通道之间的关系。1×1卷积的输出直接用作注意力权重,用于重新加权MSCA的输入。数学上,我们的MSCA可以写成:
⊗表示逐元素矩阵乘法操作。DW-Conv表示深度卷积,Scalei,i∈{0,1,2,3},表示图2(b)中的第i个分支。Scale0是恒等连接。根据[69]的方法,在每个分支中,我们使用两个深度条带卷积来近似具有大型核的标准深度卷积。在这里,每个分支的核大小分别设置为7、11和21。我们选择深度条带卷积的原因有两个。一方面,条带卷积是轻量级的。为了模拟一个具有7×7核大小的标准2D卷积,我们只需要一对7×1和1×7卷积。另一方面,在分割场景中存在一些条带状物体,例如人体和电线杆。因此,条带卷积可以作为网格卷积的补充,并有助于提取条带状特征[69, 30]。
将一系列的构建模块堆叠起来,得到了所提出的卷积编码器,称为MSCAN。对于MSCAN,我们采用了一个常见的分层结构,其中包含四个阶段,分别具有逐渐减小的空间分辨率H/4×W/4、H/8×W/8、H/16×W/16和H/32×W/32。这里,H和W分别是输入图像的高度和宽度。每个阶段包含一个下采样块和一堆构建块,如上所述。下采样块具有步长为2和核大小为3×3的卷积,然后是一个批量归一化层[36]。值得注意的是,在MSCAN的每个构建块中,我们使用批量归一化而不是层归一化,因为我们发现批量归一化对分割性能的提升更大。
我们设计了四个具有不同大小的编码器模型,分别命名为MSCAN-T、MSCAN-S、MSCAN-B和MSCAN-L。相应的整体分割模型分别称为SegNeXt-T、SegNeXt-S、SegNeXt-B和SegNeXt-L。详细的网络设置如表2所示。
表2:所提出的SegNeXt不同尺寸的详细设置。在这个表中,“e.r.”表示前馈网络中的扩张比例。‘C’和‘L’分别是通道数和构建块数。‘解码器维度’表示解码器中的MLP维度。‘参数’是在ADE20K数据集[111]上计算的。由于不同数据集中类别数目的差异,参数数量可能会略有变化。
在分割模型[90, 108, 4]中,编码器大多数是在ImageNet数据集上进行预训练的。为了捕捉高级语义,通常需要一个解码器,该解码器应用于编码器之上。在这项工作中,我们研究了三种简单的解码器结构,如图3所示。第一种结构是SegFormer [90]中采用的纯MLP结构。第二种是大多数采用CNN的模型。在这种结构中,编码器的输出直接作为输入传递给重型解码器头,如ASPP [4]、PSP [106]和DANet [20]。最后一种是我们SegNeXt中采用的结构。我们从最后三个阶段聚合特征,并使用轻量级的Hamburger [22]进一步建模全局上下文。结合我们强大的卷积编码器,我们发现使用轻量级的解码器可以提高性能-计算效率。
值得注意的是,与SegFormer的解码器从第1阶段到第4阶段聚合特征不同,我们的解码器只接收来自最后三个阶段的特征。这是因为我们的SegNeXt是基于卷积的。来自第1阶段的特征包含了太多的低级信息,对性能产生了负面影响。此外,第1阶段的操作会带来沉重的计算开销。在我们的实验部分,我们将展示我们的卷积SegNeXt的性能要比最近的基于Transformer的SegFormer [90]和HRFormer [100]要好得多。
图3:三种不同的解码器设计。
数据集。我们在七个流行的数据集上评估我们的方法,包括ImageNet-1K [14]、ADE20K [111]、Cityscapes [12]、Pascal VOC [17]、Pascal Context [65]、COCO-Stuff [3]和iSAID [84]。ImageNet [14]是最著名的图像分类数据集,包含1,000个类别。与大多数分割方法类似,我们使用它来预训练我们的MSCAN编码器。ADE20K [111]是一个具有挑战性的数据集,包含150个语义类别。它包含20,210/2,000/3,352张图像,分别用于训练、验证和测试集。Cityscapes [12]主要关注城市场景,并包含5,000张高分辨率图像,涵盖19个类别。分别有2,975/500/1,525张图像用于训练、验证和测试。Pascal VOC [17]涉及20个前景类别和一个背景类别。经过增强后,它有10,582/1,449/1,456张图像用于训练、验证和测试。Pascal Context [65]包含59个前景类别和一个背景类别。训练集和验证集分别包含4,996和5,104张图像。COCO-Stuff [3]也是一个具有挑战性的基准数据集,包含172个语义类别和总共164,000张图像。iSAID [84]是一个大规模的航拍图像分割基准数据集,包括15个前景类别和一个背景类别。它的训练、验证和测试集分别涉及1,411/458/937张图像。
实现细节。我们使用Jittor [32]和Pytorch [68]进行实验。我们的实现基于timm (Apache-2.0) [85]和mmsegmentation (Apache-2.0) [11]库,用于分类和分割。我们分割模型的所有编码器都是在ImageNet-1K数据集 [14]上进行预训练的。我们分别采用Top-1准确率和平均交并比(mIoU)作为分类和分割的评估指标。所有模型都在一个包含8个RTX 3090 GPU的节点上进行训练。
对于ImageNet的预训练,我们的数据增强方法和训练设置与DeiT [78]相同。对于分割实验,我们采用了一些常见的数据增强方法,包括随机水平翻转、随机缩放(从0.5到2)和随机裁剪。批量大小设置为8用于Cityscapes数据集,对于其他所有数据集设置为16。我们使用AdamW [61]来训练我们的模型。初始学习率设置为0.00006,并采用多项式学习率衰减策略。我们对ADE20K、Cityscapes和iSAID数据集进行了160K次迭代的训练,对于COCO-Stuff、Pascal VOC和Pascal Context数据集进行了80K次迭代的训练。在测试过程中,我们使用单尺度(SS)和多尺度(MS)翻转测试策略进行公平比较。更多细节可以在我们的补充材料中找到。
ImageNet预训练是训练分割模型的常见策略[106, 5, 90, 100, 4]。在这里,我们将我们的MSCAN与几种最近流行的基于CNN和Transformer的分类模型进行了性能比较。如表3所示,我们的MSCAN在性能上优于最近的基于CNN的方法ConvNeXt [59],并且优于流行的基于Transformer的方法,如Swin Transformer [58]和SegFormer的编码器MiT。
表5:解码器中不同注意机制的性能。SegNeXt-B w/ Ham表示MSCAN-B编码器加上Ham解码器。FLOPs是使用512×512的输入尺寸计算的。
对MSCA设计进行剔除研究。我们对ImageNet和ADE20K数据集上的MSCA设计进行了剔除研究。K × K分支包含一个深度可分离的1 × K卷积和一个K × 1深度可分离卷积。1 × 1卷积表示通道混合操作。注意力表示元素级乘积,使网络获得自适应能力。结果如表6所示。我们可以发现每个部分都对最终性能做出了贡献。
解码器中的全局上下文。解码器在分割模型中起着整合多尺度特征的全局上下文的重要作用。在这里,我们研究了不同全局上下文模块对解码器的影响。正如大多数先前的工作所示[83, 20],基于注意力的解码器在CNNs上比金字塔结构[106, 4]获得了更好的性能,因此我们只展示使用基于注意力的解码器的结果。具体而言,我们展示了4种不同类型的基于注意力的解码器的结果,包括具有O(n^2)复杂度的非局部(NL)注意力[83],以及具有O(n)复杂度的CCNet [35]、EMANet [46]和HamNet [22]。如表5所示,Ham在复杂度和性能之间实现了最佳平衡。因此,我们在解码器中使用Hamburger [22]。
解码器结构。与图像分类不同,分割模型需要高分辨率的输出。我们对分割任务进行了三种不同的解码器设计的剔除研究,所有的设计都在图3中展示。相应的结果列在表7中。我们可以看到,SegNeXt (c)获得了最佳性能,并且计算成本也较低。
我们对MSCA在分割任务中的重要性进行了实验证明。作为对比,我们按照VAN [25]的方法,将我们MSCA中的多个分支替换为一个具有大卷积核的单一卷积。如表8和表3所示,我们可以观察到虽然两个编码器在ImageNet分类任务中的性能接近,但是带有MSCA的SegNeXt比不带MSCA的设置产生了更好的结果。这表明在编码器中聚合多尺度特征对于语义分割是至关重要的。
表6:对MSCA设计的剔除研究。Top-1表示在ImageNet数据集上的Top-1准确率,mIoU表示在ADE20K基准上的mIoU。结果基于MSCAN-T模型。
表7:不同解码器结构的性能。SegNeXt-T (a)表示在解码器中使用图3(a)的结构。FLOPs是使用512×512的输入尺寸计算的。SegNeXt-T (c) w/ stage 1表示阶段1的输出也被送入解码器。
表8:我们多尺度卷积注意力(MSCA)的重要性。SegNeXt-T w/o MSCA表示我们仅使用一个具有大卷积核的分支来替代我们MSCA中的多个分支,这与[25]中的方法相同。FLOPs是使用512×512的输入尺寸计算的。
图4:SegNeXt-B和SegFormer-B2在Cityscapes数据集上的定性比较。更多的视觉结果可以在我们的附加资料中找到。
在这个小节中,我们将我们的方法与最先进的基于CNN的方法进行了比较,如HRNet [79]、ResNeSt [104]和EfficientNet [77],以及基于transformer的方法,如Swin Transformer [58]、SegFormer [90]、HRFormer [100]、MaskFormer [10]和Mask2Former [9]。
性能与计算的权衡。ADE20K和Cityscapes是语义分割中广泛使用的两个基准数据集。如图1所示,我们绘制了不同方法在Cityscapes和ADE20K验证集上的性能-计算曲线。显然,与SegFormer [90]、HRFormer [100]和MaskFormer [10]等其他最先进的方法相比,我们的方法在性能和计算之间实现了最佳的权衡。
与最先进的transformer模型的比较。我们将SegNeXt与ADE20K、Cityscapes、COCO-Stuff和Pascal Context等最先进的transformer模型进行比较。如表9所示,SegNeXt-L在ADE20K数据集上与Swin-T骨干网络的Mask2Former相比,mIoU提高了3.3个百分点(51.0对比47.7),参数和计算成本相似。此外,SegNeXt-B在ADE20K数据集上与SegFormer-B2相比,仅使用56%的计算量,mIoU提高了2.0个百分点(48.5对比46.5)。特别是,由于SegFormer [90]中的自注意力与输入尺寸呈二次复杂度的关系,而我们的方法使用卷积操作,这使得我们的方法在处理Cityscapes数据集中的高分辨率图像时表现出色。例如,SegNeXt-B在SegFormer-B2上提高了1.6个百分点的mIoU(81.0对比82.6),但使用了更少的计算量(减少了40%)。在图4中,我们还展示了与SegFormer的定性比较。我们可以看到,由于我们提出的MSCA,我们的方法在处理对象细节时能够很好地识别出来。
与最先进的CNN的比较。如表4、表10和表12所示,我们将SegNeXt与ResNeSt-269 [104]、EfficientNet-L2 [112]和HRNet-W48 [79]等最先进的CNN进行了比较,这些比较是基于Pascal VOC 2012、Pascal Context和iSAID数据集。SegNeXt-L在参数和计算量更少的情况下表现出色,超过了流行的HRNet (OCR) [79, 99]模型(60.3对比56.3),该模型专门为分割任务设计。此外,SegNeXt-L在Pascal VOC 2012测试排行榜上的表现甚至超过了在额外的3亿个不可用图像上进行预训练的EfficientNet-L2 (NAS-FPN)。值得注意的是,EfficientNet-L2 (NAS-FPN)有4.85亿参数,而SegNeXt-L只有48.7万参数。
与实时方法的比较。除了最先进的性能外,我们的方法也适用于实时部署。即使没有任何特定的软件或硬件加速,当处理尺寸为768×1,536的图像时,SegNeXt-T在单个3090 RTX GPU上实现了每秒25帧(FPS)。如表11所示,我们的方法在Cityscapes测试集上为实时分割设定了新的最先进结果。
表9:在ADE20K、Cityscapes和COCO-Stuff基准数据集上与最先进方法的比较。FLOPs(G)的数量是在ADE20K和COCO-Stuff的输入尺寸为512×512的情况下计算的,而在Cityscapes的输入尺寸为2,048×1,024的情况下计算。†表示在ImageNet-22K上进行了预训练的模型。
表10:在Pascal VOC数据集上与最先进方法的比较。∗表示使用了COCO [55]的预训练。†表示使用了JFT-300M [74]的预训练。$表示利用了额外的3亿张无标签图像进行了预训练。
表11:在Cityscapes测试数据集上与最先进的实时方法的比较。我们使用单个RTX-3090 GPU和AMD EPYC 7543 32核处理器CPU对我们的方法进行测试。在没有使用任何优化的情况下,SegNeXt-T可以实现每秒25帧(FPS),满足实时应用的要求。
在这个小节中,我们将提出的网络应用于弱监督语义分割任务。在这个任务中,通常通过使用CAM [110]由分类模型生成伪分割图。先前的工作主要利用VGGNet [72]或ResNets [28, 87]作为CAM生成器。在这里,我们测试了由我们的MSCAN生成的CAM的性能。我们使用EPS [43]架构,并按照训练策略和配方进行。数值结果如表13所示。我们可以看到,仅仅将ResNet38主干换成我们的MSCAN就可以明显提高性能,相比EPS基线。当使用我们的SegNeXt作为分割网络时,性能提升进一步增加。
表12:在Pascal Context基准测试上的比较。FLOPs的数量是以512×512的输入尺寸计算的。∗表示使用了ImageNet-22K的预训练。†表示使用了ADE20K的预训练。
表13:在PASCAL VOC 2012验证集上与之前最先进的弱监督语义分割方法的比较。除了我们使用的MSCAN-B之外,所有的分割结果都基于ResNet主干[28, 87]。
在这篇论文中,我们分析了之前成功的分割模型,并找到了它们所拥有的良好特性。基于这些发现,我们提出了定制的卷积注意力模块MSCA和CNN风格的网络SegNeXt。实验结果表明,SegNeXt在很大程度上超过了当前最先进的基于transformer的方法。 最近,基于transformer的模型在各种分割排行榜上占据主导地位。然而,本文表明,当使用适当的设计时,基于CNN的方法仍然可以比基于transformer的方法表现更好。我们希望这篇论文能够鼓励研究人员进一步探索CNN的潜力。 我们的模型也有其局限性,例如将这种方法扩展到具有1亿+参数的大规模模型以及在其他视觉或NLP任务上的性能。这些问题将在我们未来的工作中得到解决。