SETR，使用transformer+CNN做语义分割_setr网络实现

作者：Cpp五条 | 2024-04-29 11:26:59

踩

setr网络实现

摘要

语义分割目前常用的框架还是基于FCN的encoder-decoder架构。encoder的作用是提取更丰富的语义特征，一般会不断地降低特征图的大小来实现更大的感受野。感受野的大小决定了特征是否能足够捕获更大范围的周边信息甚至是全局信息，但对于语义分割，损失分辨率意味着空间损失大，分割效果可能会变差。因此文章提出用将语义分割视为序列到序列的预测任务作为替代。
文章地址：Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective
with Transformers
源码地址：SETR

目的

·目前大多数的语义分割框架采用编码器-解码器（Encoder-Decoder）结构。
编码器用于压缩原始输入图像的空间分辨率并逐步地提取更加高级的抽象语义特征。
解码器则用于将编码器所提取到的高级特征上采样到原始输入分辨率以进行像素级的预测。
上下文（context）信息是提升语义分割性能最关键的因素，而感受野（respect-field）则大致决定了网络能够利用到多少的信息。
由于网络中的有效感受野是有限的，这将会严重制约模型的表示能力。
Transformer 的一个特性便是能够保持输入和输出的空间分辨率不变，同时还能够有效的捕获全局的上下文信息。因此，作者这里便采用了类似ViT的结构来进行特征提取同时结合Decoder来恢复分辨率。

方法

SETR主要由三部分组成：输入→转换→输出。如图1。
Alt

图1 Schematic illustration of the proposed SEgmentation TRansformer (SETR) 。(a) 输入预处理及特征提取；（b）渐进式上采样；（c）多层级特征聚合。

Image to sequence

首先，需要将原始的输入图片处理成Transformer能够支持的格式，这里作者参考了ViT的做法，即对输入图像进行切片处理，将每一个2D的图像切片（patch）视为一个“1D”的序列作为整体输入到网络当中。通常来说，Transformer接收的输入是一个1维的特征嵌入序列 $Z\in R^{L\times C}$ ,其中，L为序列的长度，C为隐藏层的通道尺寸。因此，对于图像序列而言，我们也需要将输入 $x\in R^{H\times W\times 3}$ 转换为 $Z$ 。
文章采用切片的方式，每个切片大小为1616，那么对于一张256256大小的图片来说就可以切成256块。为了对每个切片的空间信息进行编码，可以为每个局部位置都学习一个特定的嵌入，并将其添加到一个线性的投影函数中来形成最终的输入序列。如此一来，尽管Transofomer是无序的，也仍然可以保留相对应的空间位置信息，因为对原始的位置信息进行了关联。

Transformer

通过将序列输入到Transformer架构可进行特征的提取，其主要包含两个部分Multi-head Self-Attention (MSA) and Multilayer Perceptron (MLP) blocks。

Decoder

关于解码器，文章给出了三种结构。经过TF所提取出的特征其输入和输出的维度是一致的，为了进行最终的像素级分割，需要将其reshape成原始的空间分辨率。

Naive upsampling (Naive)
将Transformer输出的特征维度降到分类类别数后经过双线性上采样恢复原分辨率, 也就是2-layer: 1 × 1 conv + sync batch norm (w/ ReLU) + 1 × 1 conv
Progressive UPsampling (PUP）
为了从H/16 × W/16 × 1024 恢复到H × W × 19(19是cityscape的类别数) 需要4次操作, 交替使用卷积层和两倍上采样操作来恢复到原分辨率。
Multi-Level feature Aggregation (MLA)
首先将Transformer的输出{Z1,Z2,Z3…ZLe}均匀分成M等份,每份取一个特征向量。如下图, 24个transformer的输出均分成4份,每份取最后一个,即{Z6,Z12,Z18,Z24}，后面的Decoder只处理这些取出的向量。
具体是先将ZL 从2D (H × W)/256 × C恢复到3D H/16 × W/16 × C，然后经过3-layer的卷积1 × 1, 3 × 3, and 3 × 3后再经过双线性上采样4×自上而下的融合。以增强Zl 之间的相互联系,如下图最后一个Zl理论上拥有全部上面三个feature的信息,融合，再经过3 × 3卷积后通过双线性插值4× 恢复至原分辨率。

实验

数据集
在Cityscapes[1]、ADE20K[2]以及PASCAL Context[3]这三个数据集上进行实验评估。

多尺度测试
首先将输入图像缩放到一个统一的尺寸，然后执行多尺度的缩放以及随机的水平翻转，尺度缩放因子分别为（0.5，0.75，1.0，1.25，1.5，1.75），紧接着采用滑动窗口的方式进行重叠的切片预测，最后再合并测试结果。如果移动的步长不足以得到一张完整的切片，那么以短边为例保持同等的aspect ratio。其中，由于采用的是多卡的分布式训练，因此Synchronized BN也被用于解码器和辅助损失头的训练过程中。为了简化训练过程，作者这里并没采用在线困难样本挖掘（OHEM）[4]之类的trick用于模型的训练。

模型在Cityscapes、ADE20K和PASCAL Context三个数据集上进行了实验，实验结果优于用传统FCN(with & without attention module)抽特征的方法，在ADE20k数据集上和FCN对比可视化如图2：

Alt