[论文解读] A Unified Efficient Pyramid Transformer for Semantic Segmentation

Fangrui Zhu, Yi Zhu+4 more|arXiv (Cornell University)|Jul 29, 2021

Advanced Neural Network Applications参考文献 71被引用 4

一句话总结

本文提出UN-EPT，一种统一且高效的基于Transformer的语义分割框架，通过联合优化上下文建模与边界细化，实现端到端的性能提升。通过在金字塔Transformer中采用稀疏采样策略以捕捉长距离上下文，同时引入轻量化空间分支以自适应地建模边界细节，UN-EPT在仅使用8.5GB GPU显存的情况下，于ADE20K数据集上实现了50.5的mIoU，优于参数量更大、显存消耗更高的现有方法。

ABSTRACT

Semantic segmentation is a challenging problem due to difficulties in modeling context in complex scenes and class confusions along boundaries. Most literature either focuses on context modeling or boundary refinement, which is less generalizable in open-world scenarios. In this work, we advocate a unified framework(UN-EPT) to segment objects by considering both context information and boundary artifacts. We first adapt a sparse sampling strategy to incorporate the transformer-based attention mechanism for efficient context modeling. In addition, a separate spatial branch is introduced to capture image details for boundary refinement. The whole model can be trained in an end-to-end manner. We demonstrate promising performance on three popular benchmarks for semantic segmentation with low memory footprint. Code will be released soon.

研究动机与目标

• 为解决现有语义分割方法在上下文建模与边界细化之间过度偏重其一而导致在开放世界场景中性能受限的问题。
• 开发一种统一框架，同时提升复杂场景中全局上下文理解能力与对象边界细化能力。
• 降低标准Transformer在密集预测任务（如语义分割）中带来的高显存与高计算成本。
• 实现端到端训练，有效结合多尺度上下文特征与细粒度空间细节。
• 在极低显存占用下实现SOTA性能，且无需大规模预训练。

提出的方法

• 在基于Transformer的解码器中采用稀疏采样策略，以减少注意力计算量与显存占用，实现对高分辨率特征图的高效处理。
• 提出一种高效金字塔Transformer（EPT）模块，利用主干网络提取的多尺度特征，增强不同感受野下的上下文表征能力。
• 设计独立的空间分支，包含三层主干网络与两个头，用于捕捉局部图像细节并优化分割边界。
• 将空间分支的特征作为输入自适应查询输入到Transformer解码器中，实现基于局部图像内容的动态上下文建模。
• 将上下文分支（EPT）与空间分支在端到端可训练的统一网络（UN-EPT）中联合优化，实现分割质量与边界精度的协同提升。
• 采用集合预测范式，将输入设为展平后的图像像素序列，输出为类别标签序列，支持基于Transformer的序列到序列学习。

实验结果

研究问题

• RQ1 是否能够设计一种统一的深度学习框架，在不牺牲效率的前提下，有效平衡语义分割中的上下文建模与边界细化？
• RQ2 如何使自注意力机制在计算上足够高效，以处理密集预测任务中的高分辨率输入？
• RQ3 多尺度特征融合与稀疏注意力对复杂场景中长距离依赖建模有何影响？
• RQ4 轻量化空间分支在不增加模型复杂度的前提下，能在多大程度上提升边界定位精度？
• RQ5 与现有方法相比，基于Transformer的架构是否能在显著降低显存消耗的同时实现SOTA性能？

主要发现

• 当使用DeiT作为主干网络时，UN-EPT在ADE20K基准上实现了50.5的mIoU，优于采用相同架构的先前方法。
• 当使用ResNet50作为主干网络时，UN-EPT在ADE20K上实现了46.1的mIoU，表明其在无需大规模预训练的情况下仍具备强大性能。
• 模型训练期间仅需8.5GB GPU显存，显著低于SETR（30.0GB）及其他基于Transformer的模型。
• 消融实验表明，空间分支使mIoU提升1.1个百分点，验证了其在边界细化中的有效性。
• 在金字塔各层级中，每张特征图使用16个采样点时性能达到最优，进一步增加采样点数量对性能影响微乎其微。
• UN-EPT在所有对比方法中GFLOPs最低（99.1），表明其具有极高的计算效率。