【论文学习】基于 Transformer 的图像分割模型

基于 Transformer 的图像分割模型

SegViT: Semantic Segmentation with Plain Vision Transformers

SegViT：基于纯视觉 Transformer 的语义分割

理解

Vision Transformer (ViT)：将Transformer模型创新性地应用到了计算机视觉领域

具体做法：

图像分块：将一张图片切割成多个固定大小的小图像块。
线性嵌入：将这些图像块展平成向量序列。
送入Transformer：把这个向量序列（加上位置编码）像处理句子中的单词序列一样，送入标准的Transformer编码器进行学习

在ViT之前，主流的视觉模型都是卷积神经网络（CNN） 。很多研究试图将Transformer和CNN结合，或者对Transformer进行大量针对视觉任务的定制化修改。

而这里的 "纯" 指的是SegViT这项工作直接使用最原始、最标准的Transformer编码器（即ViT的核心）作为其骨干网络，而没有引入CNN结构或进行非常复杂的、针对分割任务的架构

摘要

我们探索了朴素视觉Transformer（ViT）在语义分割任务中的应用能力，并提出了SegViT模型。以往的基于ViT的分割网络通常从ViT的输出中学习像素级表征。与之不同，我们利用其核心组件------注意力机制来生成语义分割的掩码。具体而言，我们提出了注意力到掩码（ATM）模块，通过计算一组可学习的类别令牌与空间特征图之间的相似度图谱，并将其转换为分割掩码。实验表明，采用ATM模块的SegViT在ADE20K数据集上优于使用朴素ViT骨干网络的同类模型，并在COCO-Stuff-10K和PASCAL-Context数据集上达到了新的最优性能。此外，为降低ViT骨干网络的计算成本，我们提出基于查询的下采样（QD）和基于查询的上采样（QU）方法，构建了一种收缩结构。采用该收缩结构后，模型在保持竞争力的性能同时可节省高达40%的计算量。代码已通过以下链接开源：https://github.com/zbwxp/SegVit。

理解1

以往的基于ViT的分割网络通常从ViT的输出中学习像素级表征。

一个标准的ViT模型（如ViT-B/16）处理图像的过程是：

图像分块：将图像切成16x16像素的块。
Transformer处理：这些图像块经过Transformer编码器处理。
输出：Transformer编码器的输出包括：
1. 一个[CLS]令牌的输出：通常用于图像分类，代表了整个图像的全局信息。
2. 一系列图像块令牌的输出：每个令牌代表了原始图像中一个16x16像素块的信息。

ViT的原始输出是一系列图像块令牌，它已经丢失了原始图像中的二维空间结构。对于需要输出一张与原图大小对应的分割图的语义分割任务来说，这是不够的。

学习像素级表征：设计一个额外的"解码器"模块，其任务是将ViT输出的、失去空间结构的图像块序列，重新"翻译"和"上采样"回一个密集的像素级表征（即分割图）

这个过程通常包括：

恢复空间结构：首先，将ViT输出的图像块令牌序列，根据它们在原始图像中的位置，重新排列成一个二维的特征图。这个特征图的分辨率会很低（例如，输入图像尺寸的1/16）。
逐步上采样：使用各种上采样技术（如转置卷积、双线性插值等）来逐步提高这个低分辨率特征图的空间尺寸。
融合多层特征：为了弥补上采样过程中丢失的细节，通常会从ViT的不同深度（层）提取特征。浅层特征包含更多空间细节（如边缘、纹理），深层特征包含更多语义信息（如"这是辆车"）。解码器会将这些特征融合起来。
预测：最后，这个经过上采样和融合得到的高分辨率、像素级表征会通过一个分类层，为每个像素分配一个类别标签。

理解2

我们提出了注意力到掩码（ATM）模块，通过计算一组可学习的类别令牌与空间特征图之间的相似度图谱，并将其转换为分割掩码。

输入：

一组可学习的类别令牌 ：例如，模型有150个类别要分割，它就有150个可学习的令牌。每个令牌都是一个数值向量。在训练过程中，这些向量会通过梯度下降自动学习，最终每个令牌都编码了对应类别的核心视觉特征。
空间特征图：是从ViT主干网络中提取出的、保留了二维空间结构的特征图。特征图上的每一个点都有一个特征向量，这个向量描述了原始图像对应区域的内容。

计算"相似度图谱"：相似度计算 ：ATM模块会计算每一个"类别令牌" 与空间特征图上的每一个位置的特征向量 之间的相似度。对于一个类别令牌 （比如"天空"），计算完成后，我们会得到一张二维的图。这张图上每个像素的值，代表了"天空"这个类别与原始图像对应区域的匹配程度。------这个过程为每一个类别都生成了一张专属的相似度图谱。

从相似度图谱到分割掩码：对于空间特征图上的每一个位置，我们找出哪个类别的相似度图谱在该位置的值最高 ；直接将一组相似度图谱转换成了一张最终的、每个像素都有唯一类别标签的分割掩码

理解3

此外，为降低ViT骨干网络的计算成本，我们提出基于查询的下采样（QD）和基于查询的上采样（QU）方法，构建了一种收缩结构。

Vision Transformer (ViT) 有一个致命缺点：其自注意力机制的计算量随着输入序列长度的增加呈平方级增长。图像被切成的patch越多（即序列越长），计算就越慢、显存占用越高

基于查询的下采样（QD）：

目标：减少序列长度（即降低特征图的空间尺寸，例如从H x W降到H/2 x W/2）
传统做法：使用池化或卷积 stride=2。这种方法简单、固定，但可能丢失重要信息
QD的做法 ：
1. 引入可学习的查询 ：定义一组数量较少（对应下采样后的序列长度）的、可学习的向量，我们称之为查询令牌。
2. 交叉注意力 ：让这些"查询令牌"作为查询，去"询问"当前所有的图像块令牌（键和值）。这个过程就是交叉注意力
3. 产出：通过注意力加权求和，每个"查询令牌"都从原始的所有图像块令牌中聚合了全局信息，生成了一个新的、更具代表性的令牌
4. 结果：我们用这组数量更少但信息密度更高的新令牌，替代原来庞大的令牌序列。序列长度成功减小，但关键信息被保留了下来

基于查询的上采样（QU）：

目标：增加序列长度（即提高特征图的空间尺寸，例如从H/2 x W/2恢复到H x W）。
传统做法：双线性插值或转置卷积。这些方法是固定的或内容无关的，恢复的细节可能不够精确。
QU的做法 ：
1. 引入位置查询 ：定义一组数量较多（对应上采样后的序列长度）的向量。这些向量不仅可学习，还包含了目标高分辨率网格的位置信息 （即位置编码）。我们称之为位置查询。
2. 交叉注意力 ：让这些"位置查询"作为查询，去"询问"当前低分辨率特征图的令牌（键和值）。
3. 产出：每个"位置查询"根据自己所在的目标位置，从低分辨率特征中提取最相关的信息，生成一个新的、对应高分辨率位置的令牌。
4. 结果：我们得到了一组数量更多、空间分辨率更高的令牌序列。细节信息根据内容被智能地恢复和填充。

Masked-attention Mask Transformer for Universal Image Segmentation

基于掩码注意力机制的掩码Transformer通用图像分割模型

理解

将分割任务转化为 "掩码 - 类别" 联合预测，将所有分割任务统一为 "生成一组掩码，并为每个掩码预测类别" 的问题

对语义分割：每个类别对应一个掩码；对实例分割：每个实例对应一个掩码；对全景分割：同时覆盖语义（背景类别）和实例（前景物体）的掩码集合

掩码注意力：让 Transformer 专注于每个掩码内部的像素关联，同时抑制不同掩码间的干扰

具体工作流程如下：

初始化可学习的 "查询向量（Query）"：每个查询向量对应一个待预测的掩码（数量根据任务动态设定，如实例分割通常设 100 个查询以覆盖潜在实例）。
交叉注意力提取全局特征：查询向量与图像特征（由 CNN 或 Transformer 骨干提取）通过交叉注意力交互，初步学习每个查询对应的区域特征。
掩码注意力细化掩码边界：
1. 基于当前查询生成的 "初始掩码"，计算掩码内像素的注意力权重（仅关注掩码覆盖的像素区域），忽略掩码外的背景或其他物体像素。
2. 这种 "掩码内聚焦" 机制避免了不同物体区域的特征混淆，让模型更精准地捕捉单个掩码的细节（如物体边缘、孔洞）。
迭代优化：通过多轮 Transformer 解码，逐步细化掩码形状和类别预测，最终输出每个查询对应的 "掩码 + 类别" 对。

摘要

图像分割旨在将具有不同语义（如类别或实例隶属关系）的像素进行分组。每种语义定义对应一个任务。尽管各任务间仅语义定义不同，但当前研究却集中于为每个任务设计专用架构。我们提出了掩码注意力掩码Transformer（Mask2Former），这是一种能够处理任何图像分割任务（全景、实例或语义）的新架构。其关键组件包括掩码注意力机制，该机制通过将交叉注意力限制在预测的掩码区域内来提取局部特征。该架构不仅将研究工作量减少了至少三倍，更在四个主流数据集上以显著优势超越了现有最佳专用架构的性能。尤其值得注意的是，Mask2Former在全景分割（COCO数据集上57.8 PQ）、实例分割（COCO数据集上50.1 AP）和语义分割（ADE20K数据集上57.7 mIoU）任务上均创造了新的性能纪录。