论文地址:https://arxiv.org/abs/2105.15203
1 引言
文章提出了一种基于transformer的语义分割网络,不同于ViT模型,SegFormer使用一种分层特征表示的方法,每个transformer层的输出特征尺寸逐层递减,通过这种方式捕获不同尺度的特征信息。并且舍弃了ViT中的position embedding操作,避免了测试图像与训练图像尺寸不同而导致模型性能下降的问题。在decoder部分采用简单的MLP结构,聚合transformer层不同尺度的特征,可以同时融合局部注意力和全局注意力。
SegFormer-B4以64M参数在ADE20K上达到50.3% mIoU,比之前最好的模型小5倍,准确率高2.2%。SegFormer-B5在Cityscapes验证集上达到84.0% mIoU。
2 模型
SegFormer由两个部分组成:
- 一个分层的Transformer encoder,用来提取大尺度的粗粒度特征和小尺度的细粒度特征。
- 一个轻量级的ALL-MLP decoder,用来混合encoder中不同层输出的特征并产生最终的语义分割遮罩。
2.1 Hierarchical Transformer Encoder
2.1.1 Hierarchical Feature Representation
不同于ViT中产生的单一尺寸的特征映射,文中的Hierarchical Transformer Encoder会产生像是CNN网络一样的多层次特征,靠近输入的层提供大尺度的粗粒度信息,靠近输出的层提供小尺度的细粒度信息,通过聚合不同尺度的特征表示可以有效的提升模型在语义分割任务中的表现。
2.1.2 Overlapped Patch Merging
文中通过Overlapped Patch Merging来实现特征图尺寸递减的操作,减少特征图的尺寸而增加特征图的通道数。这种操作的弊端是不能保留不同patch之间的连续信息。
2.1.3 Efficient Self-Attention
由于Attention操作的复杂度与序列长度的平方成正比,文中采用一种序列减少操作(先reshape在通过一个线性层)降低序列长度,以此来减少计算的复杂度。
2.1.4 Mix-FFN
ViT模型中的positional encoding操作不适用于训练图像和测试图像尺寸不同的情况,文章提出一种Mix-FFN结构,通过添加一个3×3卷积和一个MLP层到FFN中,实验证明3×3卷积足够给Transformer层提供位置信息。
2.2 Lightweight ALL-MLP Decoder
SegFormer的decoder层仅由MLP构成,它主要包含四步:
- 将transformer各层的输出转换成统一的额维度。
- 将特征上采样然后进行拼接。
- 使用MLP混合拼接的不同特征。
- 最后使用一个MLP生成遮罩。
2.3 Relationship to SERT
SegFormer与SERT的对比:
- SegFormer使用的预训练数据集更小。
- SegFormer使用层次化的encoder,参数量更小,且能同时捕获不同尺度的特征。
- SegFormer摒弃了Positional Embedding的操作,可以适应训练集和测试集图像尺寸不同的情况。
- SegFormer更加紧凑,且需要更少的计算资源。
3 总结
文章中包含大量的设计细节,且更改了原始的transformer结构,使得模型计算量更小,虽然整体的设计思路并无太大亮点,但在小细节上下了很多功夫,值得参考。