Segmenter: Transformer for Semantic Segmentation

论文：[2105.05633] Segmenter: Transformer for Semantic Segmentation (arxiv.org)

代码：[rstrudel/segmenter: ICCV2021] Official PyTorch implementation of Segmenter: Transformer for Semantic Segmentation (github.com)

期刊/会议：ICCV 2021

摘要

图像分割通常在单个图像patch的级别上是模糊的，并且需要上下文信息来达成标签共识。本文介绍了一种用于语义分割的transformer模型——segmenter。与基于卷积的方法相比，我们的方法允许在第一层和整个网络中建模全局上下文。我们以最近的Vision Transformer(ViT)为基础，并将其扩展到语义分割。为此，我们依赖于与图像patch对应的输出嵌入，并使用逐点(point-wise)线性解码器或mask transformer解码器从这些嵌入中获取类标签。我们利用预训练的图像分类模型，并表明我们可以在用于语义分割的中等大小数据集上对它们进行微调。线性解码器已经可以获得很好的结果，但是可以通过mask transformer产生类掩码来进一步提高性能。我们进行了广泛的消融研究，以显示不同参数的影响，特别是大模型和小patch尺寸的性能更好。Segmenter在语义分割中取得了很好的效果。它在ADE20K和Pascal上下文数据集上的性能都优于目前的SOTA，在Cityscape上也具有竞争力。

1、简介

语义分割是一个具有挑战性的计算机视觉问题，具有广泛的应用，包括自动驾驶、机器人、增强现实、图像编辑、医学成像等。语义分割的目标是将每个图像像素分配给对应于底层对象的类别标签，并为目标任务提供高级图像表示，例如为虚拟试穿应用检测人的边界和他们的衣服。尽管近年来进行了大量的努力并取得了很大的进展，但由于丰富的类内变化、上下文变化和由遮挡和低图像分辨率引起的模糊性，图像分割仍然是一个具有挑战性的问题。

最近的语义分割方法通常依赖于卷积编码器-解码器架构，其中编码器生成低分辨率的图像特征，解码器对特征进行上采样，并使用逐像素类得分进行分割映射。最先进的方法部署了全卷积网络(FCN)，并在具有挑战性的分割基准上取得了令人印象深刻的结果。这些方法依赖于可学习的堆叠卷积，可以捕获语义丰富的信息，并在计算机视觉中非常成功。然而，卷积滤波器的局域特性，限制对图像全局信息的访问。同时，这些信息对于分割尤其重要，因为局部patch的标注往往依赖于全局图像上下文。为了避免这个问题，DeepLab方法引入了空洞卷积和空间金字塔池化的特征聚合。这允许扩大卷积网络的感受野，并获得多尺度特征。根据NLP的最新进展，一些分割方法探索了基于通道或空间注意力和点注意力的替代聚合方案，以更好地捕获上下文信息。然而，这样的方法仍然依赖于卷积骨干网络，因此偏向于局部交互。广泛使用专门的层来纠正这种偏置表明了卷积架构在分割方面的局限性。

为了克服这些限制，我们将语义分割问题表述为一个seq2seq的问题，并使用transformer架构在模型的每个阶段利用上下文信息。根据设计，transformer可以捕捉场景元素之间的全局交互，并且没有内置的感应先验，参见图1。然而，全局交互建模的成本是二次的，这使得这种方法在应用于原始图像像素时非常昂贵。根据最近在Vision Transformer(ViT)上的工作，我们将图像分割为patch，并将线性补丁嵌入作为Transformer编码器的输入token。然后，由编码器产生的上下文化token序列由Transformer解码器上采样到逐像素类分数。对于解码，我们考虑将patch嵌入简单地按点线性映射到类分数或基于Transformer的解码方案，其中可学习的类嵌入与patch token一起处理以生成类掩码。我们通过消融实验对模型正则化、模型大小、输入patch大小及其在精度和性能之间的权衡，对transformer进行了扩展研究。我们的Segmenter方法在保持简单、灵活和快速的同时获得了出色的结果。特别是，当使用具有小输入patch大小的大型模型时，最佳模型在具有挑战性的ADE20K数据集上达到了53.63%的平均IoU，大大超过了之前所有最先进的卷积方法5.3%。这种改进部分源于我们的方法在模型的每个阶段捕获的全局上下文，如图1所示。

总之，我们的工作提供了以下四个贡献:(i)我们提出了一种基于Vision Transformer(ViT)的新的语义分割方法，它不使用卷积，通过设计捕获上下文信息，并且优于基于FCN的方法。(ii)我们提出了一系列具有不同分辨率的模型，允许在精度和运行时间之间进行权衡，从最先进的性能到具有快速推断和良好性能的模型。(iii)我们提出了一种基于transformer的解码器，它生成的类掩码优于我们的线性基线，并且可以扩展到执行更一般的图像分割任务。(iv)我们证明，我们的方法在ADE20K和Pascal Context数据集上都产生了SOTA，并且在Cityscape上具有竞争力。

2、相关工作

语义分割：基于全卷积网络(FCN)与编码器-解码器结构相结合的方法已成为语义分割的主流方法。最初的方法依赖于连续堆叠的卷积，然后是空间池化来执行密集预测。连续方法对高级特征图进行采样，并在解码时将其与低级特征图结合起来，以捕获全局信息并恢复锐利的物体边界。为了扩大第一层卷积的感受野，有几种方法提出了空洞卷积或可分离卷积。为了在更高层次上捕获全局信息，最近的工作使用空间金字塔池来捕获多尺度的上下文信息。Deeplabv3+将这些增强功能与可分离空间金字塔池化结合起来，提出了一种简单有效的编码器-解码器FCN架构。最近的工作用编码器特征映射上的注意力机制取代粗粒度池化，以更好地捕获长期依赖关系。

虽然目前的分割方法大多集中在改进FCN，但**卷积对局部操作的限制可能意味着全局图像上下文处理效率低，分割结果不佳。**因此，我们提出了一个纯transformer架构，在编码和解码阶段在模型的每一层捕获全局上下文。

Vision Transformer：在许多自然语言处理(NLP)任务中，Transformer是目前最先进的方法。这种模型依赖于自注意机制，并捕捉句子中token(单词)之间的长期依赖关系。此外，transformer非常适合并行化，便于在大型数据集上进行训练。transformer在NLP中的成功启发了计算机视觉中的几种方法，将CNN与自注意机制相结合，以解决目标检测，语义分割，全景分割，视频处理和少样本分类。

最近，Vision Transformer(ViT)引入了一种用于图像分类的无卷积 transformer架构，其中输入图像被处理为patch token序列。虽然ViT需要在非常大的数据集上进行训练，但DeiT提出了一种基于token的蒸馏策略，并在ImageNet-1K数据集上使用CNN作为训练教师模型，获得了具有竞争力的Vision Transformer模型。并发工作将这项工作扩展到视频分类和语义分割。更详细地说，SETR使用ViT骨干网络和标准CNN解码器。Swin Transformer使用ViT的变体，由局部窗口组成，在层和UpperNet之间移动，作为金字塔FCN解码器。

在这里，我们提出了Segmenter，一种用于语义图像分割的Transformer编码器解码器架构。我们的方法依赖于ViT骨干网络，并引入了一个受DETR启发的掩码解码器。我们的架构不使用卷积，通过设计捕获全局图像上下文，并在标准图像分割基准测试中获得具有竞争力的性能。

3、我们的方法：Segmenter

Segmenter是完全基于transformer的编码器-解码器架构，将一系列patch嵌入映射到像素级类标注。该模型的概述如图2所示。patch序列由3.1节中描述的transformer编码器编码，并由3.2节中描述的逐点线性映射或掩码transformer解码。我们的模型采用逐像素交叉熵损失进行端到端训练。在推理时，上采样后应用argmax以获得每个像素的单个类。

3.1 编码器

3.2 解码器

我们的mask transformer的灵感来自于DETR， MaxDeepLab和SOLO-v2，它们引入了对象嵌入来产生实例掩码。然而，与我们的方法不同，MaxDeepLab使用了一种基于CNN和transformer的混合方法，并由于计算限制将像素和类嵌入分成两个流。使用纯Transformer架构和利用patch级编码，我们提出了一种在解码阶段联合处理patch和类嵌入的简单方法。这种方法允许产生随输入变化的动态滤波器。当我们在这项工作中处理语义分割时，我们的掩码transformer也可以通过用对象嵌入取代类嵌入来直接进行泛视分割。

4、实验结果

4.1 数据集和评估指标

数据集：ADE20K、Pascal Context、CityScape。

评估指标：MIOU。

4.2 实施细节

Transformer模型：对于编码器，我们在Vision Transformer ViT的基础上构建，并考虑表1中描述的“Tiny”、“Small”、“Base”和“Large”模型。Transformer编码器中变化的参数是层数和token大小。multi-headed self-attention (MSA) block的head大小固定为64，head数量是token大小除以head大小，MSA后面MLP的隐藏大小是token大小的4倍。我们也使用DeiT，一个ViT的变体。我们考虑代表不同分辨率图像的模型，并使用输入patch大小8×8, 16×16和32 × 32。在下文中，我们使用缩写来描述模型变体和patch大小，例如SegB/16表示输入patch大小为16 × 16的“Base”变量。基于DeiT的模型用†表示，例如Seg-B†/16。

ImageNet上预训练：我们的Segmenter模型在ImageNet上进行预训练，ViT在ImageNet-21k上进行预训练，具有强大的数据增强和正则化，其变体DeiT在ImageNet-1k上进行预训练。原始的ViT模型仅使用随机裁剪进行训练，而有些工作提出的训练程序使用dropout和随机深度的组合作为正则化，Mixup和RandAugment作为数据增强。这显著提高了ImageNet top-1的精度，即在ViT-B/16上获得了+2%的增益。我们在ADE20K上微调了ViT-B/16。接下来，所有的Segmenter模型都将使用改进的ViT模型进行初始化。我们使用图像分类库timm和谷歌research提供的公开模型。这两个模型都在图像分辨率为224的情况下进行了预训练，并在ImageNet-1k上以384的分辨率进行了微调，除了ViT-B/8已以224的分辨率进行了微调。我们保持patch大小固定，并根据数据集以更高的分辨率微调语义分割任务的模型。由于patch大小是固定的，提高分辨率会导致更长的token序列。我们根据预先训练好的位置嵌入在图像中的原始位置进行双线性插值，以匹配微调序列长度。3.2节中描述的解码器初始化时使用截断正态分布的随机权重。

数据增强。在训练过程中，我们遵循语义分割库MMSegmentation的标准管道处理，这意味着裁剪，随机调整图像大小到0.5和2.0之间的比例，以及随机左右翻转。我们随机裁剪大图像，并将小图像填充为固定大小，ADE20K为512×512, Pascal-Context为480×480, Cityscape为768×768。在ADE20K上，我们训练了我们最大的模型Seg-L-Mask/16，分辨率为640×640，与Swin Transformer使用的分辨率相匹配。

优化器。为了对语义分割任务的预训练模型进行微调，我们使用了标准的像素级交叉熵损失，而不进行权重再平衡。我们使用随机梯度下降(SGD)作为基础学习率γ 0的优化器，并将权重衰减设置为0。在DeepLab的开创性工作之后，我们采用了“poly”学习率衰减γ = γ 0 ( 1 − N _{i t e r} / N _{t o t a l} ) ^0.9 ，其中N _{i t e r} 和 N _{t o t a l} 表示当前迭代数和总迭代数。对于ADE20K，我们将基础学习率γ 0 设置为10−3，并训练160K次迭代，batch size为8。对于Pascal Context，我们将γ 0设置为10−3，并训练80K迭代，batch size为16。对于Cityscape，我们将γ 0设置为10−2，并训练80K迭代，batch size为8。该计划类似于DeepLabv3+，学习速率除以因子10，除了Cityscape，我们使用因子1。

推理。为了在推断过程中处理不同的图像大小，我们使用了一个与训练大小匹配的分辨率的滑动窗口。对于多尺度推断，遵循标准实践，我们使用缩放因子为(0.5,0.75,1.0,1.25,1.5,1.75)的图像的重新缩放版本，并对结果进行左右翻转和平均。

4.3 消融实验

在本节中，我们将在ADE20K验证集上对比我们方法的不同变体。我们研究了模型正则化、模型大小、patch大小、模型性能、训练数据集大小，将Segmenter与卷积方法进行比较，并评估不同的解码器。除非另有说明，我们使用基线线性解码器并使用单尺度推断报告结果。

正则化：我们首先比较了两种形式的正则化，dropout和随机深度，并表明随机深度始终改善transformer分割训练效果。CNN模型也依赖批归一化BN，也是一种正则化。而Transformer在训练时通常由层归一化LN结合dropout作为正则化器组成。Dropout随机忽略作为块输入的token，随机深度在正向传递过程中随机跳过模型的可学习块。我们比较了基于ViT-S/16骨干网络的Seg-S/16上的正则化。从表2可以看出，将随机深度设置为0.1，随机降低10%的层数，性能持续提高，当dropout设置为0时，与未正则化的基线相比，提高了0.36%。正则化始终会影响性能，无论是单独还是与随机深度结合。这与一些工作中观察到dropout对图像分类的负面影响是一致的。从现在开始，所有的模型都将随机深度设置为0.1，并且没有dropout。

Transformer大小：我们现在研究Transformer尺寸对性能的影响，方法是在固定patch大小为16的情况下改变层数和token大小。表3显示了性能与骨干网络大小的良好关系。当token维度翻倍时，从Seg-S/16到Seg-B/16，我们得到了2.69%的改进。当层数翻倍时，从Seg-B/16到Seg-L/16，我们得到了2.65%的改进。最后，我们最大的Segmenter模型Seg-L/16在ADE20K验证数据集上使用单一尺度推理的简单解码方案实现了50.71%的mIoU。FCN模型中大量使用的任务特定层的缺失表明基于transformer的方法提供了更有表现力的模型，非常适合于语义分割。

patch大小：用patch序列表示图像提供了一种通过改变patch大小在速度和准确性之间权衡的简单方法。虽然增加patch大小会导致图像的表现更粗糙，但是在更小尺寸的序列中它将处理的更快。表3的第三和第四部分报告了ViT作为骨干网络和不同patch大小的性能。我们观察到，patch大小是影响语义分割性能的一个关键因素。这对模型大小同样重要。事实上，从32到16的patch大小，我们观察到Seg-B改善了5%。对于Seg-B，我们也报告了patch大小为8的结果，并报告了mIoU为49.54%，将ViT-B/8到ViT-L/16的差距缩小到1.17%，同时需要的参数大大减少。这一趋势表明，减少patch大小是一种稳健的改进来源，它不引入任何参数，但需要在较长的序列上计算注意力，增加了计算时间和内存占用。如果计算上可行，ViT-L/8可能是性能最好的模型。更高效计算和内存的transformer处理更大序列的小patch是一个有前景的方向。

为了进一步研究patch大小的影响，我们在图3中显示了由Segmenter模型生成的随着patch大小减小的分割图。我们观察到，对于32的patch大小，模型学习了一个全局有意义的分割，但产生了较差的边界，例如，左边的两个人是由一个blob预测的。当观察人体轮廓时可以观察到，减小patch大小会导致相当清晰的边界。很难分割实例，因为背景中的细街灯杆仅以8的分辨率捕获。在表4中，我们报告了关于对象大小的平均IoU，并将Segmenter与DeepLabv3+与ResNeSt骨干网络进行了比较。为了重现DeepLabv3+结果，我们使用了MMSegmentation库中的模型。我们观察到，与Seg-B/16相比，Seg-B/8的改进主要来自小型和中型实例，分别获得1.27%和1.74%的增益。此外，我们观察到Segmenter相对于DeepLab的最大改进来自于大型实例，其中Seg-L-Mask/16显示了6.39%的改进。

不同的解码器：在本节中，我们将比较不同的解码器变体。我们评估3.2节中介绍的mask transformer，并将其与线性基线进行比较。mask transformer有两层，具有与编码器相同的token和隐藏大小。表4报告了mIOU性能。mask transformer在线性基线上提供一致的改进。Seg-B†/16获得了最显著的1.6%的增益，Seg-B-Mask/32获得了1.1%的改进，Seg-L/16获得了0.6%的增益。在表4中，我们还检查了不同模型的增益与目标大小的关系。我们观察了小型和大型对象的增益，显示了使用动态滤波器(dynamical filter)的好处。在大多数情况下，大对象的增益更为显著，例如，Seg-B/32为1.4%，Seg-B†/16为2.1%，Seg-L/16为1.7%。mask transformer学习的类嵌入在语义上是有意义的，也就是说，相似的类就在附近，参见图8了解更多细节。

Transformer对比FCN：表4和表6比较了我们的方法与FCN模型以及DeepLabv3+与ResNeSt骨干网络，这是最好的全卷积方法之一。我们的transformer方法提供了对这种最先进的卷积方法的显著改进，突出了transformer捕捉全局场景理解的能力。segmenter在大型实例上的表现始终优于DeepLab, Seg-L/16的改进超过4%，Seg-L-Mask/16的改进超过6%。然而，DeepLab在小型和中型实例上的性能与Seg-B/16相似，同时具有相似的参数数量。Seg-B/8和Seg-L/16在小型和中型实例上表现最好，但计算成本较高。

效果：在图4中，我们将我们的模型与几种最先进的方法进行了比较，就每秒处理图像数量和mIoU而言，并显示了Segmenter优于基于FCN的模型的明显优势(绿色曲线)。我们还表明，我们的方法与最近的基于transformer的方法相比，我们最大的模型Seg-L-Mask/16与Swin-L相当，并且优于SETR-MLA。我们观察到Seg/16模型在精度与计算时间方面表现最好，而Seg- B-Mask/16提供了很好的权衡。Seg-B-Mask/16在推理速度相似的情况下优于基于FCN的方法，与SETR-MLA匹配，但速度快两倍，所需参数更少，在推理速度和性能方面优于Swin-B。Seg/32模型学习上一节讨论的更粗粒度的分割映射，并为Seg-B-Mask /32提供每秒400张图像的快速推断，会比ResNet-50快4倍，但是性能差不多。为了计算每秒的图像，我们使用V100 GPU，将图像分辨率固定为512，对于每个模型，我们最大限度地提高内存允许的batch size，以便进行公平的比较。

数据集的大小：Vision Transformer强调了大数据集在图像分类任务中获得良好性能的重要性。在语义分割数据集的规模上，我们分析了表5中使用越来越大的数据集训练时Seg-S/16在ADE20k数据集上的性能。当训练集大小低于8k图像时，我们观察到性能的重要下降。这表明，即使在微调过程中，transformer在足够的数据量下表现最好。

4.4 对比SOTA

在本节中，我们将在ADE20K、Pascal上下文和城市景观数据集上比较分段器的性能与最先进的方法。

ADE20K。在ImageNet-1k上预训练的Seg-B†/16与最先进的FCN方法DeepLabv3+ ResNeSt200[63]相匹配，如表6所示。加入我们的掩码变压器，Seg-B†-Mask/16提高了2%，达到50.08% mIoU，优于所有FCN方法。我们的最佳模型Seg-L-Mask/16达到了53.63%的最先进性能，超过了mIoU DeepLabv3+ ResNeSt-200和基于变压器的方法SETR[67]和Swin-L UperNet[35]的5.27%。

Pascal Context性能如表7所示。Seg-B†模型与FCN方法具有竞争力，更大的Seg-L/16模型已经提供了最先进的性能，优于SETR-L。我们的掩码变压器Seg-L-Mask/16可以进一步提高性能，比线性解码器提高2.5%，达到59.04% mIoU的性能。特别地，我们报告了Segmenter比OCR HRNetV2-W48提高2.8%，比SETR-L MLA提高3.2%。

Cityscapes：表8报告了segmenter在Cityscape上的性能。我们使用了Seg-L-Mask/16的mask transformer的变体，在解码器中只有一层，因为由于768×768的大输入分辨率，两层不适合内存。Seg-B和Seg-L方法与其他最先进的方法相比都具有竞争力，Seg-L-Mask/16的mIoU达到81.3%。

定性的结果。图5显示了segmenter和DeepLabv3+与ResNeSt骨干网络的定性比较，其中模型由MMSegmentation库提供。我们可以观察到Deeplabv3+倾向于生成更清晰的对象边界，而Segmenter在大型实例上提供更一致的标签，并更好地处理部分遮挡。

5、总结

本文介绍了一种纯transformer的语义分割方法。编码器部分建立在最近的Vision transformer(ViT)上，但不同之处在于我们依赖于所有图像patch的编码。我们观察到transformer很好地捕获了全局上下文。将简单的逐点线性解码器应用于patch编码已经取得了很好的效果。用mask transformer解码进一步提高了性能。我们相信，我们的端到端编码器-解码器transformer是朝着语义分割、实例分割和全景分割的统一方法迈出的第一步。