原文链接:Per-Pixel Classification is Not All You Need for Semantic Segment
代码:MaskFormer

问题定义和动机

自 FCN 诞生以来,语义分割任务就被默认当作逐像素分类的任务来解决。诚然,这种方法简化了语义分割任务,但是它仍然存在很多问题,其中最大的问题就是这种方法只能输出固定个数的分割掩码,所以逐像素分类很难解决实例分割这样的更难的任务。但是反观实例分割,这种任务一直以来都是用以 Mask R-CNN 为代表的掩码分类的方法来解决的。掩码分类和逐像素分类的区别就在于掩码分类会预测一组二值掩码,将每一个掩码与单个类别关联,而不是每个像素都需要类别。因此,在这篇论文里,作者提出了以下两个问题:

  1. 能否找到一个简洁通用的 mask classification 模型同时解决语义分割和实例分割问题?
  2. 这个 mask classification 模型在语义分割上的结果能否超越传统的像素分类模型

论文方法

MaskFormer结构

MaskFormer的基本思想是把语义分割、实例分割、甚至全景分割都用一个统一的框架、损失和训练过程来实现。MaskFormer与 DETR 思想类似,都是去除 NMS/anchor 等人工先验,把实例分割作为一个集合预测问题来看待,具体的网络结构如图 1 所示,整体包含了三个模块。像素级模块的作用是生成像素级特征嵌入,它先用 backbone 提取图像特征得到低分辨率的特征图,然后通过 pixel decoder 模块进行逐步上采样生成和原图宽高一致的像素级特征嵌入 $\varepsilon_{pixel}$ 。Transformer 模块主要包含了一个标准的 Transformer 解码器,解码器的输入有两部分,一部分是来自像素级模块的特征图,另一部分是 N 个可学习的位置嵌入 (N queries),这里其实是受 DETR 的启发,用 object query 去进行后续的掩码预测和类别预测。在分割模块中,Transformer 模块的输出 Q 会经过一个线性层并用 softmax 函数激活得到类别的概率预测,这里图中的 +1 指的是不包含任何类别的部分。掩码预测是通过将 Q 经过 MLP 后与 $\varepsilon_{pixel}$ 点乘并用 sigmoid 函数激活得到的。

MaskFormer结构

损失函数

本文中作者借鉴了 DETR 中 bipartite matching loss 的设计思想来训练模型,所用的损失函数为:

$$
\mathcal{L}_{\text{mask-cls}}(z, z^{\text{gt}}) = \sum_{j=1}^{N} \left[ -\log p_{\sigma(j)}(c_{j}^{\text{gt}}) + \Bbb{1}_{c_{j}^{\text{gt}} \neq \varnothing} \mathcal{L}_{\text{mask}}(m_{\sigma(j)}, m_{j}^{\text{gt}}) \right]
$$

其中 $-\log p_{\sigma(j)}(c_{j}^{\text{gt}})$ 是分类损失, $\Bbb{1}_{c_{j}^{\text{gt}}\ne \varnothing}\mathcal{L}_{\text{mask}}(m_{\sigma(j)},m_{j}^{\text{gt}})$ 是掩码损失,$\mathcal{L}_{\text{mask}}$ 是二进制掩码损失。大多数现有的掩码分类模型还使用辅助损失,然而本论文提出的简单的掩码分类模型,允许单独使用掩码分类损失 $\mathcal{L}_{\text{mask-cls}}$ 进行端到端训练。

推理过程

论文中提出了两种推理过程,一种是通用的推理过程,另一种是专门用于语义分割的推理过程。作者认为推理策略的选择在很大程度上取决于评估指标,而不是任务本身。

  1. 通用推理
    通过将每个像素点 $[h, w]$ 分配给 $N$ 个预测的概率-掩码对中的一个来将图像划分为段。直观上,只有当最可能的类别概率和掩码预测概率都很高时,才会将位置 $[h, w]$ 的像素分配给相应的概率-掩码对 $i$。因此分配公式为 $\arg\max_{i:c_i \neq \varnothing} p_i(c_i) \cdot m_i[h, w]$ ,其中 $c_i$ 是最可能的类别标签。被分配到相同概率-掩码对 i 的像素形成一个段,其中每个像素都被标记为 $c_i$。对于语义分割,类别标签相同的段会被合并;而对于实例级分割任务,使用概率-掩码对的索引 $i$ 区分同一类别的不同实例。最后,为了降低全景分割中的误报,作者在推理前过滤掉了低置信度的预测,并移除了那些预测掩码($m_i > 0.5$)的大部分被其他预测遮挡的预测段。

  2. 语义推理
    语义推理通过一个简单的矩阵乘法来实现。作者通过实验发现,对概率-掩码对进行边缘化,即 $\arg\max_{c=1,\ldots,K} \sum_{i=1}^{N} p_i(c) \cdot m_i[h, w]$ ,比通用推理中使用的策略得到的结果更好。$\arg\max$ 不包括“无对象”类别($\varnothing$),因为标准的语义分割要求每个输出像素都要有一个标签。

论文贡献

本文中作者提出了基于掩码分类的语义分割模型 MaskFormer,并通过实验证明了掩码分类足够通用,甚至可以使用完全相同的模型、损失和训练过程,以统一的方式解决语义级和实例级分割任务,为语义分割任务提供了新的解决思路
论文中 MaskFormer 在语义分割中表现出色,在多个语义分割数据集上,都优于逐像素分类的方法。尤其是在数据集包含的类别数量较多的情况下,MaskFormer 的表现更好。在全景分割任务中它也表现优异,优于复杂的 DETR 模型和 Max-DeepLab 模型,并且无需复杂的辅助损失,简化了模型设计。本论文通过实验验证了掩码分类相对于逐像素分类在分割任务中的优势,为后续研究提供了新的思路和方向。