标题
Per-Pixel Classification is Not All You Need for Semantic Segmentation (NIPS 2021)
论文:Per-Pixel Classification is Not All You Need for Semantic Segmentation
代码:https://github.com/facebookresearch/MaskFormer
摘要
现代方法通常将语义分割视为逐像素分类任务,而实例级分割则通过掩码分类的替代方案进行处理。我们的核心观点是:掩码分类具有足够强的通用性,能够使用完全相同的模型、损失函数和训练流程,以统一的方式同时解决语义级和实例级分割任务。基于这一观察,我们提出了MaskFormer——一个简单的掩码分类模型,其通过预测一组二值掩码(每个掩码对应一个全局类别标签预测)来实现分割任务。总体而言,这种基于掩码分类的新方法简化了语义分割与全景分割任务的有效解决方案体系,并展现出卓越的实证效果。特别地,我们发现当类别数量较大时,MaskFormer的表现显著优于传统的逐像素分类基线方法。我们基于掩码分类的方法在语义分割(ADE20K数据集55.6 mIoU)和全景分割(COCO数据集52.7 PQ)任务上均超越了当前最先进的模型。
1 引言

语义分割的目标是将图像划分为具有不同语义类别的区域。自Long等人[30]提出全卷积网络(FCN)以来,大多数基于深度学习的语义分割方法都将该任务建模为逐像素分类(图1左),即对每个输出像素施加分类损失函数[9,52]。在这种范式下,逐像素的预测结果自然会将图像划分为不同类别的区域。
掩码分类是一种替代范式,它将图像划分与分类这两个分割任务的核心要素进行解耦。与逐个像素分类不同,基于掩码分类的方法通过预测一组二值掩码来实现分割,其中每个掩码对应一个全局类别标签预测(图1右)。这种更灵活的掩码分类范式在实例级分割领域占据主导地位。无论是Mask R-CNN[21]还是DETR[4],在实例分割和全景分割任务中都采用每个分割区域对应单一类别预测的机制。相比之下,逐像素分类方法假设输出数量固定,无法返回实例级任务所需的动态数量预测区域/片段。
我们的核心观察发现:掩码分类具有足够的通用性,能够同时解决语义级和实例级分割任务。事实上,在FCN[30]提出之前,性能最佳的语义分割方法(如O2P[5]和SDS[20])就采用了掩码分类范式。基于这一视角,自然引发出两个关键问题:单一掩码分类模型能否简化语义级和实例级分割任务的有效方法体系?同时,这样的掩码分类模型在语义分割任务中能否超越现有的逐像素分类方法?
为解答上述问题,我们提出了一种简洁的 MaskFormer 方法,其能够将现有任意逐像素分类模型无缝转换为掩码分类范式。通过采用DETR[4]提出的集合预测机制,MaskFormer利用Transformer解码器[41]计算一组预测对,其中每个预测对包含一个类别预测和一个掩码嵌入向量。该掩码嵌入向量通过与底层全卷积网络生成的逐像素嵌入进行点积运算,最终生成二值掩码预测。这一新型模型以统一方式同时解决语义级和实例级分割任务:无需修改模型架构、损失函数或训练流程。具体而言,无论是语义分割还是全景分割任务,MaskFormer均采用相同的逐像素二值掩码损失函数与每个掩码对应的单一分类损失进行监督训练。最后,我们设计了简洁的推理策略,将模型输出适配到不同任务所需的预测格式。
我们在五个具有不同类别数量的语义分割数据集上对MaskFormer进行评估:Cityscapes[15](19类)、Mapillary Vistas[34](65类)、ADE20K[55](150类)、COCOStuff-10K[3](171类)以及ADE20K-Full[55](847类)。实验表明,在类别较少的Cityscapes数据集上,MaskFormer与逐像素分类模型性能相当;而在类别规模更大的数据集上,新模型展现出显著优势。我们推测,每个掩码对应单一类别预测的机制相比逐像素分类,能更好地建模细粒度识别任务。当采用Swin-Transformer[29]作为骨干网络时,MaskFormer在ADE20K数据集上取得55.6 mIoU的最新最佳性能,较同骨干的逐像素分类模型[29]提升2.1 mIoU,同时具有更高效率(参数量减少10%,计算量降低40%)。
最后,我们通过COCO[28,24]和ADE20K[55]两个全景分割数据集验证MaskFormer处理实例级任务的能力。MaskFormer在使用相同骨干网络和后处理流程的条件下,性能优于更复杂的DETR模型[4]。此外,该模型在COCO数据集上以52.7 PQ刷新全景分割性能纪录,较先前最佳方法[42]提升1.6 PQ。系列实验充分验证了MaskFormer统一实例分割与语义分割任务的技术优势。
2 相关工作
逐像素分类与掩码分类方法在语义分割领域均得到了深入研究。早期工作中,Konishi和Yuille[25]基于局部图像统计特征应用逐像素贝叶斯分类器。随后受非语义分组研究的启发[13,36],基于掩码分类的方法逐渐兴起,并在PASCAL VOC挑战赛[18]中取得最佳性能。O2P[5]和CFM[16]等方法通过对候选掩膜[6,40,2]进行分类,实现了当时最先进的分割效果。2015年,FCN[30]将逐像素分类思想引入深度网络,在mIoU指标(该指标特别契合逐像素分类的分割范式)上显著超越所有先前方法。
自全卷积网络(FCN)[30]这一开创性工作提出以来,逐像素分类成为基于深度学习的语义分割主流范式。现代语义分割模型主要聚焦于在最终特征图中聚合长程上下文信息:ASPP[7,8]采用不同膨胀率的空洞卷积;PPM[52]使用多种核尺寸的池化算子;DANet[19]、OCNet[51]和CCNet[23]则运用非局部模块[43]的多种变体。近期,SETR[53]与Segmenter[37]采用视觉Transformer(ViT)[17]替代传统卷积骨干网络,从网络底层即开始捕获长程上下文。值得注意的是,这些基于Transformer[41]的语义分割方法仍沿用逐像素分类范式。需特别说明的是,本文提出的MaskFormer模块可将任意逐像素分类模型转换为掩码分类范式,从而无缝集成该领域的技术进展。
掩码分类范式在实例级分割任务中应用广泛[20,24]。此类任务需要动态数量的预测输出,而逐像素分类方法假设输出数量固定,难以满足这一需求。占据主导地位的Mask R-CNN[21]采用全局分类器对实例分割的候选掩膜进行分类。DETR[4]进一步引入Transformer[41]架构,实现对全景分割[24]中物体(thing)与背景(stuff)的同步分割。然而,这些掩码分类方法需要预测边界框,可能限制其在语义分割中的应用。最新提出的Max-DeepLab[42]通过条件卷积[39,44]技术消除了对边界框预测的依赖,但除主体掩码分类损失外,仍需引入多个辅助损失(包括实例区分损失、掩码ID交叉熵损失以及标准逐像素分类损失)。
3 从逐像素到掩码分类
本小节中,我们首先阐述如何将语义分割问题形式化为逐像素分类或掩码分类两种范式。随后,我们将介绍如何基于Transformer解码器[41]构建掩码分类模型的具体实现方案。最后,我们将说明如何通过简洁的推理策略将掩码分类输出转换为适应不同任务需求的预测格式。
3.1 逐像素分类
在逐像素分类中,语义分割模型的目标是为 \(H\)×\(W\) 尺寸图像中的每个像素预测其在 \(K\) 个可能类别上的概率分布:\(y = \{p_i | p_i \in \Delta^K\}^{H \cdot W}_{i=1}\),此处 \(\Delta^K\) 表示 \(K\) 维概率单纯形(probability simplex)。训练逐像素分类模型的过程直观简洁:给定每个像素的真实类别标签 \(y^{gt} = \{y^{gt}_i | y^{gt}_i \in \{1,...,K\}\}^{H \cdot W}_{i=1}\),通常采用逐像素交叉熵(负对数似然)损失函数进行优化,即 \(\mathcal{L}_{pixel-cls}(y, y^{gt}) = \sum_{i=1}^{H·W} -\text{log}\ p_i(y^{gt}_i)\)。
3.2 掩码分类
掩码分类将分割任务分解为两个步骤:
1)将图像划分为 \(N\) 个区域(\(N\) 无需等于类别数 \(K\)),每个区域由二值掩码 \(\{m_i | m_i \in [0,1]^{H×W}\}^N_{i=1}\)表示;
2)将每个区域视为整体关联至 \(K\) 个类别的概率分布。
为实现区域划分与分类的联合建模(即执行掩码分类),我们将模型输出 \(z\) 定义为包含 \(N\) 个概率-掩码对的集合,即 \(z = \{(p_i, m_i)\}^N_{i=1}\)。与逐像素分类不同,掩码分类中的概率分布 \(p_i \in \Delta^{K+1}\) 在 \(K\) 个类别标签基础上增加了辅助的"无对象(no object)"标签( \(\varnothing\) )。该 \(\varnothing\) 标签用于标识未对应任何类别的掩码预测。值得注意的是,掩码分类允许多个掩码预测关联至同一类别,因此可同时适用于语义级与实例级分割任务(应该是指一个掩码除了可以对应一个类别外,还可以对应一个实例)。
训练掩码分类模型时,需在预测集合 \(z\) 与真实标注分割区域集合 \(z^{gt} = \{(c^{gt}_i, m^{gt}_i) | c^{gt}_i \in \{1,...,K\}, m^{gt}_i \in \{0, 1\}^{H \times W} \}^{N^{gt}}_{i=1}\)之间建立匹配关系 \(\sigma\)。此处 \(c^{gt}_i\)表示第 \(i\) 个真实区域的类别标签。由于预测集合大小 \(|z|=N\) 与真实集合大小 \(|z^{gt}|=N^{gt}\) 通常不等,我们设定 \(N ≥ N^{gt}\) 并通过填充"无对象"标签 \(\varnothing\) 使二者可进行一对一的匹配。
对于语义分割任务,当预测数量 \(N\) 等于类别数 \(K\) 时可采用固定匹配策略(fixed matching):第 \(i\) 个预测直接对应类别标签 \(i\) 的真实区域(若该类别不存在于真实标注中,则匹配至 \(\varnothing\))。实验表明,基于二分匹配(bipartite matching-based)的分配策略相比固定匹配能获得更优性能。与DETR[4]采用边界框计算预测 \(z_i\) 与真实 \(z^{gt}_j\) 之间的匹配代价不同,我们直接基于类别预测概率和掩码损失函数计算匹配代价,即:\(-p_i(c^{gt}_j) + \mathcal{L}_{\text{mask}}(m_i, m^{gt}_j)\),其中 \(\mathcal{L}_{\text{mask}}\) 为二值掩码损失函数。
在模型参数训练过程中,给定匹配关系 \(\sigma\) 后,主要掩码分类损失 \(\mathcal{L}_{\text{mask-cls}}\) 由每个预测区域的交叉熵分类损失与二值掩码损失 \(\mathcal{L}_{\text{mask}}\) 共同构成:
值得注意的是,现有大多数掩码分类模型除 \(\mathcal{L}_{\text{mask-cls}}\) 外还需引入辅助损失(如边界框损失[21,4]或实例区分损失[42])。在下一节中,我们将提出一种简洁的掩码分类模型,其仅需 \(\mathcal{L}_{\text{mask-cls}}\) 即可实现端到端训练。
3.3 MaskFormer

现在我们介绍新型掩码分类模型——MaskFormer,其通过计算 \(N\) 个概率-掩码对 \(z = \{(p_i, m_i)\}_{i=1}^N\)实现分割任务。该模型包含三大模块(见图2):
1) 像素级模块,负责提取用于生成二值掩码预测的逐像素嵌入(per-pixel embeddings);
2) Transformer模块,通过堆叠Transformer解码层[41]计算 \(N\) 个逐片段嵌入(per-segment embeddings);
3) 分割模块,基于上述嵌入生成预测对 \(\{(p_i, m_i)\}_{i=1}^N\)。推理过程中(详见3.4节),\(p_i\) 与 \(m_i\) 将被组合为最终预测结果。

像素级模块以 \(H\times W\) 尺寸图像作为输入。骨干网络首先生成低分辨率特征图 \(\mathcal{F} \in \mathbb{R}^{C_\mathcal{F} \times (\frac{H}{S})\times(\frac{W}{S})}\) ,其中 \(C_\mathcal{F}\) 为通道数,\(S\) 为特征图下采样步长(\(C_\mathcal{F}\) 取决于具体骨干网络,本工作中采用 \(S=32\))。随后,像素解码器通过逐步上采样操作生成逐像素嵌入 \(\mathcal{E}_{\text{pixel}} \in R^{C_\mathcal{E}\times H\times W}\),\(C_\mathcal{E}\) 为嵌入维度。值得注意的是,任何基于逐像素分类的分割模型(包括近期基于Transformer的模型[37,53,29])均可适配此像素级模块设计。MaskFormer能够无缝地将此类模型转换为掩码分类范式。

Transformer模块采用标准Transformer解码器[41],基于图像特征 \(\mathcal{F}\) 和 \(N\) 个可学习的位置嵌入(即查询向量)并行输出 \(N\) 个逐片段嵌入\(Q \in \mathbb{R}^{C_\mathcal{Q} \times N}\)(\(C_\mathcal{Q}\) 为嵌入维度),这些嵌入向量编码了MaskFormer预测的每个分割区域的全局信息。与DETR[4]类似,该解码器可实现所有预测结果的并行生成。

分割模块在逐片段嵌入 \(Q\) 的基础上应用线性分类器与softmax激活函数,生成每个区域的类别概率预测 \(\{p_i \in \Delta^{K+1}\}^N_{i=1}\)。值得注意的是,该分类器额外预测"无对象"类别(\(\varnothing\)),用于标识未对应任何区域的嵌入向量。对于掩码预测,采用包含两个隐藏层的多层感知机(MLP)将逐片段嵌入 \(Q\) 转换为 \(N\) 个维度为 \(C_\mathcal{E}\) 的掩码嵌入 \(\mathcal{E}_\text{mask} \in \mathbb{R}^{C_\mathcal{E}\times N}\)。最终,通过将第 \(i\) 个掩码嵌入与像素级模块生成的逐像素嵌入 \(\mathcal{E}_\text{pixel}\) 进行点积运算,并施加sigmoid激活函数,即可获得二值掩码预测 \(m_i \in [0,1]^{H\times W}\)。具体而言,\(m_i[h,w] = \text{sigmoid}(\mathcal{E}_\text{mask}[:,i]^\text{T} \cdot \mathcal{E}_\text{pixel}[:,h,w])\)。
需要强调的是,通过实验我们发现:不使用softmax激活函数强制掩码预测互斥能够获得更好的性能。在训练过程中,\(\mathcal{L}_\text{mask-cls}\) 损失函数将交叉熵分类损失与二值掩码损失 \(\mathcal{L}_\text{mask}\) 相结合。为保持简洁,我们采用与DETR[4]相同的 \(\mathcal{L}_\text{mask}\) 定义,即分别用超参数 \(\lambda_\text{focal}\) 和 \(\lambda_\text{dice}\) 加权后的 focal损失[27] 与 dice损失[33] 的线性组合。
3.4 掩码分类推理
首先,我们提出一种通用的推理流程,可将掩码分类输出 \(\{(p_i, m_i)\}_{i=1}^N\) 转换为适用于全景分割或语义分割的输出格式。随后,我们将描述专为语义分割设计的语义推理策略。需要指出的是,具体推理策略的选择主要取决于评估指标而非任务类型。
通用推理流程通过将每个像素 \([h,w]\) 分配至满足 \(\arg \max_{i:c_i\neq \varnothing} p_i(c_i) \cdot m_i[h,w]\) 条件的概率-掩码对 \(i\),实现图像区域划分。其中 \(c_i\) 表示每个概率-掩码对 \(i\) 的最可能类别标签 \(c_i = \arg \max_{c \in \{1,...,K,\varnothing\}} p_i(c)\)。直观而言,该策略仅在像素位置 \([h,w]\) 对应的类别概率 \(p_i(c_i)\) 与掩码预测概率 \(m_i[h,w]\) 均较高时,才会将该像素分配至第 \(i\) 个概率-掩码对。被分配至同一概率-掩码对 \(i\) 的像素构成一个分割区域,其类别标签为 \(c_i\)。对于语义分割任务,同类别的区域将被合并;而对于实例级分割任务,概率-掩码对的索引 \(i\) 可用于区分同类别不同实例。最后,为降低全景分割中的误检率,我们沿用先前的推理策略[4,24]:在推理前过滤低置信度预测,并移除二值掩码中大部分区域(\(m_i > 0.5\))被其他预测遮挡的分割区域。
语义推理策略专为语义分割设计,通过简单的矩阵乘法实现。实验表明,对概率-掩码对进行边缘化处理(即 \(\arg \max_{c\in \{1,...,K\}} \sum_{i=1}^N p_i(c) \cdot m_i[h,w]\))相比通用策略中的硬性像素分配能获得更优性能。由于标准语义分割要求每个像素必须具有类别标签,\(\arg\max\) 运算中不包含"无对象"类别(\(\varnothing\))。值得注意的是,该策略实际上返回逐像素类别概率 \(\sum_{i=1}^N p_i(c) \cdot m_i[h,w]\)。然而我们发现,直接最大化逐像素类别似然会导致性能下降。推测原因是梯度被平均分配至每个查询向量,从而增加了训练难度。