SeMask: Semantically Masked Transformers for Semantic Segmentation（论文解读）

这是一篇2023发布在ICCV期刊上论文

一、motivation 

论文：ICCV 2023 Open Access Repository (thecvf.com)

以往的基于Transformer的语义分割网络一般都是对一个backbone进行微调达到Encoder的效果，在Decoder中才会考虑上下文的信息进行特征融合；这种方法忽略了图像在编码阶段提供的语义上下文信息。文章提出SeMask有效框架，借助语义注意力操作将语义信息整合到编码器中。并且，在训练过程中使用轻量级语义解码器，在每个阶段为中间语义先验图提供监督。

代码：Picsart-AI-Research/SeMask-Segmentation: [NIVT Workshop @ ICCV 2023] SeMask: Semantically Masked Transformers for Semantic Segmentation (github.com)

二、相关工作

缺乏语义上下文信息会导致次优分割性能，特别是在小对象的情况下，这些对象会与较大类别的边界合并，从而导致错误的预测。

backbone：分层视觉Transformer

图 1：流行的基于Transformer的分割网络（左）和 SeMask（右）之间的比较。与大多数现有的直接使用预训练主干而不进行任何更改的方法相比，SeMask 通过添加一个额外的语义层在编码器主干中使用语义先验；这种简单的更改显着提高了性能。

三、总体架构

原理：主干的每个阶段的Transformer层之后插入一个语义层（SeMask）；论文使用轻量级语义解码器来累积来自所有阶段的语义图，并使用标准解码器（如semantic FPN）来进行主要的逐像素预测。在整个编码器中添加的语义建模和特征建模有助于我们提高语义分割任务的性能。

四、Encoder

大小为H×W×3的RGB输入图像首先被分割成大小为4×4的非重叠块。较小大小的补丁支持分割中的密集预测。这些补丁充当令牌，并作为分层视觉Transformer编码器的输入。

由四个阶段组成，每个阶段都包含了NA个Swin Transformer的Transformer层和NS个SeMask注意块的语义层。每个阶段的特征图为原始图像分辨率的{1/4,1/8,1/16,1/32}，K为类别数，C为通道数

经过Encoder之后输入到轻量级语义解码器的是语义先验图，输入到FPN的是语义掩蔽特征

Transformer层

组成：由NA个Swin Transformer块组成；用以从图像中提取图像级上下文信息

公式如下：

其中加入了一个相对位置嵌入(RPE) ，以此考虑到元素之间的位置关系

网络图如下：得到的Y作为SeMask语义模块的输入

其中对于分层表示，我们通过下一阶段的补丁合并层将特征图从 H/4 × W/4 缩小到 H/8 × W/8。对下一阶段迭代该补丁合并以获得分层特征图，分辨率为，其中 i ∈ {1, 2, 3, 4}。 

 语义层

组成：包含NS个SeMask Attention 块；将语义信息与特征解耦，产生语义先验图（每个像素属于不同语义类别的概率分布），在这些语义先验映射的指导下更新特征。

公式如下：

网络图如下：

𝑌′是通过语义掩蔽的自注意力机制加强的特征，将原始特征与语义信息相结合；不仅基于原始的特征，还结合了通过语义掩码调整的查询所提供的额外语义信息。

而 𝑆是语义先验图，提供了直接关于数据点类别的信息，可用于后续的分类或分割任务。 

五、Decoder

轻量级语义解码器：聚合语义先验特征，以在训练期间预测网络的语义先验；在编码器的每个阶段为语义先验图提供基本事实监督，提高了语义分割任务的性能。

语义FPN解码器：聚合语义掩蔽特征；通过一系列卷积、双线性上采样和求和运算融合了不同阶段的特征，以产生最终的密集像素预测。使其成为我们目的的分割解码器，既高效又直观。

六、损失函数

总损失LT计算为每像素两个交叉熵损失的总和：L1(主损失函数)和L2辅助损失函数;

辅助损失函数：有助于在训练过程中提供额外的监督信号;

其中：

• F：网络的主要预测;
• S：语义先验预测;
• §：表示将存储在 GT （Ground Truth）中实况类标签转换为 one-hot 格式（正确类别的位置是1，其余位置是0）;
• Lce ：是交叉熵损失,我们根据经验设置 

通过结合L1和 L2：SeMask框架确保了模型不仅在最终输出层能做出准确的预测，同时在整个网络的各个层次都能学习到有效的特征表示

七、实验

实验部分只进行了总结，具体的细节可看原论文。

数据集：

• ADE20K。ADE20K是一个覆盖150个细粒度语义概念的场景解析数据集，是最具挑战性的语义分割数据集之一。该训练集包含20210幅具有150个语义类的图像。验证和测试集分别包含2000张和3352张图像。
• 城市景观。 Cityscapes是一个用于语义分割的城市街道驱动数据集，由来自50个城市的5000张图像组成，共有19个语义类。训练集中有2975幅图像，验证集中有500幅图像，测试集中有1525幅图像。

Transformer 模块：

• Tin和Small变体在ImageNet-1k上进行训练，分辨率为224×224。
• †代表ImageNet-22k在384×384分辨率图像上的预训练。 

网络初始化：

• 用公开的模型进行初始化
• 解码器使用正态分布的随机权重进行初始化

数据扩充：

• 在训练过程中，我们执行平均减法，将图像缩放到从（0.5，0.75，1.0，1.25，1.5，1.75）随机采样的比例，随机左右翻转和颜色抖动。
• 我们随机裁剪大图像，并将小图像填充到固定大小，ADE20K为512×512，Cityscapes为768×768。
• 在ADE20K上，我们以640×640的分辨率训练我们最大的模型Semask-L†FPN，与SwinTransformer使用的分辨率相匹配。

Transformer 尺寸消融实验：

在cityscape数据集上的改进3个百分点，因为分割数据集中的类更少，创建了更强的先验。

Semantic Attention消融实验：

 通过用一个简单的单头自注意块代替SeMask块来证明它的有效性，这会损害Tiny变体的性能。这说明简单的注意并不能提高语义注意的有效性和有效性

可学习常数（λ）消融实验：

λ的包含对SeMask块的成功是有效的，因为它充当了修改特征的调谐因子，从而控制了权重初始化的噪声。实验证明，在编码器中不同阶段的推理过程中，λ∈[0.05，0.3]。 

辅助损失消融实验：

 Query Swin-T FPN使用来自Transformer层的查询进行损耗计算。可以观察到SeMask表现最好。

结论： 

使用SeMask swin - L†作为编码器，Mask2Former-MSFaPN作为主要预测的解码器，我们在单尺度和多尺度mIoU度量上分别获得了57.00%和58.25%的分数。 

八、代码问题

代码：Picsart-AI-Research/SeMask-Segmentation: [NIVT Workshop @ ICCV 2023] SeMask: Semantically Masked Transformers for Semantic Segmentation (github.com)

作者给的环境：

我跑的是Semask-FPN，用的数据是ADE20K，没有使用预训练参数

用作者给的版本运行得过程中遇到太多问题了，前前后后改了有两三天，以下是我自己使用的版本：

pytorch和cuda的版本：Previous PyTorch Versions | PyTorch

conda install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge

mmcv的版本：安装 MMCV — mmcv 1.7.0 文档 

pip install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.10/index.html

pytorch和cuda、mmcv三者的关系一定要对应，去官网看相应的对应版本

 就是这三者的版本对应比较折磨人，因为我用低版本的pytorch和cuda版本不行，所有升了它们的版本，进而也升了mmcv的版本，升了这个版本的话，以下代码就需要注释掉，去掉mmcv的版本检查：

其他的话按照论文作者的过程没什么问题了