点击下方卡片，关注「集智书童」公众号
自从Vision Transformers问世以来，许多计算机视觉任务（如语义分割）的格局最近发生了重大变革，而这些任务一直由CNNs主导。然而，计算成本和内存需求使得这些方法不适用于移动设备，特别是对于高分辨率的每像素语义分割任务。
在本文中介绍了一种用于移动语义分割的squeeze-enhanced Axial Transormer（SeaFormer）新方法。具体来说，设计了一个以squeeze Axial和detail enhancement为特征的通用注意力块。它可以进一步用于创建具有优异成本效益的主干架构系列。结合轻量化分割头，在ADE20K和Cityscapes数据集上实现了基于ARM的移动设备上分割精度和延迟之间的最佳权衡。
至关重要的是，本文方法以更好的性能和更低的延迟击败了 mobile-friendly对手和基于Transformer的对手。除了语义分割，作者还将所提出的SeaFormer架构应用于图像分类问题，展示了作为通用mobile-friendly骨干的潜力。
1、简介
语义分割是计算机视觉中的一个基本问题，其目的是为图像中的每个像素分配一个语义类标签。传统的方法依赖于叠加局部卷积核来感知图像的长距离结构信息。
自从Vision Transformers问世以来，语义分割领域发生了重大变革。基于Transformer的方法已经显著地展示了全局上下文建模的能力。然而，Transformer的计算成本和内存需求使得这些方法不适用于移动设备，尤其是高分辨率图像输入。
遵循高效操作的传统智慧，引入了基于局部/窗口的注意力、轴向注意力、动态图消息传递和一些轻量级注意力机制。
然而，由于高分辨率输入的高延迟，这些进步仍然不足以满足移动设备的设计要求和约束（见图1）。最近，人们对构建基于Transformer的语义分割产生了极大的兴趣。为了在高分辨率下降低计算成本，TopFormer致力于在原始输入的1/64尺度上应用全局注意力，但这肯定会损害分割性能。
为了以性能无害的方式解决像素级分割任务的高分辨率计算和移动设备上的低延迟要求的困境，作者提出了一系列基于移动友好Transformer的语义分割模型，称为SeaFormer，它将轴向注意力的计算复杂性从O((H+W)HW)降低到O(HW)，以在移动设备上实现卓越的精度-效率权衡，并填补移动友好型高效Transformer的空缺。
核心构建块为squeeze-enhanced Axial attention（SEA注意力）试图将输入特征图沿水平/垂直轴挤压（合并）成紧凑的列/行，并计算自注意力。通过concat query、key和value以补偿在压缩过程中牺牲的细节信息，然后将其馈送到深度卷积层以增强局部细节。
再加上一个轻量化分割头，在小规模特征中采用了拟议的SeaFormer层，能够在移动设备上进行低延迟的高分辨率图像语义分割。如图1所示，所提出的SeaFormer在ADE20K数据集上以较低的延迟优于其他有效的神经网络。特别是，SeaFormer Base在基于ARM的移动设备上具有较低的延迟（106ms vs.126ms），优于轻量级的CNN对手MobileNetV3（41.0 vs.33.1mIoU）。
主要贡献：
- 引入了一种用于移动语义分割的squeeze-enhanced Axial Transformer（SeaFormer）框架；
- 设计了一个以 squeeze Axial 和 detail enhancement公式为特征的通用注意力块；它可用于创建一系列具有卓越成本效益的主干架构；
- 在ADE20K和Cityscapes数据集上表现出色，以明显的优势击败了移动友好竞争对手和基于Transformer的细分模型；
- 除了语义分割之外，还将所提出的SeaFormer架构应用于图像分类问题，证明了作为通用移动友好骨干的潜力。
2、本文方法
2.1、整体架构
受这两种分支架构的启发，作者设计了一种squeeze-enhanced Axial Transformer（SeaFormer）框架。如图2所示，SeaFormer由以下部分组成：共享STEM、上下文分支、空间分支、融合块和轻量化分割头。为了公平比较，遵循TopFormer设计STEM。它由一个stride为2的卷积和4个MobileNet block组成，其中第一个和第三个块的stride是2。上下文分支和空间分支共享生成的特征图，这允许构建快速的语义分割模型。
1、上下文分支
上下文分支被设计为从特征图中捕获上下文丰富的信息。如图2的红色分支所示，上下文分支分为3个阶段。为了获得更大的感受野，在将MobileNet block应用于下采样和扩展特征维度之后堆叠SeaFormer层。与作为下采样模块的标准卷积相比，MobileNet block增加了模型的表示能力，同时保持了较低的计算量和延迟。
对于除SeaFormer Large以外的变体，SeaFormer层应用于最后两个阶段，以在精度和效率之间实现卓越的权衡。对于SeaFormer Large，在上下文分支的每个阶段插入SeaFormer层。为了在分割精度和推理速度之间实现良好的权衡，设计了一个SEA注意力。
2、空间分支
空间分支被设计为获得高分辨率的空间信息。与上下文分支相同，空间分支重用特征映射。然而，早期卷积层的特征包含丰富的空间细节，但缺乏高级语义信息。因此，设计了一个融合块，将上下文分支中的特征融合到空间分支中，将高级语义信息引入到低级空间信息中。
3、融合Block
如图2所示，空间分支中的高分辨率特征图之后是1×1卷积和批量归一化层，以生成要融合的特征。上下文分支中的低分辨率特征图被馈送到1×1卷积层、批归一化层、sigmoid层，并被上采样到高分辨率，以通过双线性插值产生语义权重。
然后，上下文分支的语义权重被元素乘以空间分支的高分辨率特征。融合块使得低级空间特征能够获得高级语义信息。
4、轻量化分割Head
最后一个融合块后的特征被直接输入所提出的分割头，如图2所示。
为了快速推理，轻量化分割头由2个卷积层组成，然后分别是批归一化层，并将批归一化层的特征输入一个激活层。
2.2、 SQUEEZE-ENHANCED轴向注意力
全局注意力可以表示为
其中，。q，k，v是x的线性投影，例如，，，，其中，是可学习的权值。表示位置的特征图上的所有位置。
当将传统的注意力模块应用于H×W×C的特征图上时，时间复杂度可以为，导致效率低，延迟高。
为了提高效率，在局部区域内有一些计算自注意力的工作。在等式2、3中展示了两个最具代表性的高效Transformer。
方程2由基于窗口的注意力表示，成功地将时间复杂度降低到，其中表示O的相邻m×m位置，但失去了全局感受野。
方程3由轴向注意力表示，其仅将时间复杂度降低到，其中表示O列的所有位置；表示o行的所有位置。
根据它们的缺点，作者提出了一种移动友好的 squeeze-enhanced Axial 注意力，它具有用于全局语义提取的 squeeze Axial attention和用于 local details supplement的高效卷积细节增强内核。
1、Squeeze Axial attention
为了实现更高效的计算并同时聚合全局信息，采用了更激进的策略。同样，q，k，v首先从x得到，。根据等式4首先通过在水平方向上获取query特征图的平均值来实现水平挤压。同样，右侧显示垂直方向上的垂直挤压。z→(·)表示按给定排列张量z的维数，并且是所有元素都等于1的向量。对q的挤压操作也在k和v上重复，因此最终得到。挤压操作将全局信息保留到单个轴，从而大大减轻了等式5所示的以下全局语义提取。
特征图的每个位置只在两个压缩的轴向特征上传播信息。虽然与公式3相比，它没有明显的计算减少，但公式5的重复可以简单地通过最有效的广播操作来实现。具体情况如图3所示。压缩q，k，v的时间复杂度为，而注意力操作则需要时间。因此，挤压轴向注意力成功地降低了时间复杂度到。
2、Squeeze Axial position embedding
然而，式4不具有位置感知性，不包括特征映射的位置信息。因此，提出了挤压轴向位置嵌入来挤压轴向注意力。为了挤压轴向注意力，作者通过引入位置嵌入，从可学习参数B线性插值，，来了解它们在压缩轴向特征中的位置。L是一个常数。同样，将应用于、。因此，位置感知挤压轴向注意力可以表示为式6。
3、Detail enhancement kernel
挤压操作虽然能有效地提取全局语义信息，但却牺牲了局部细节。因此，采用了一种基于卷积的辅助来增强空间细节。如图3所示的上路径，q，k，v首先从x与另一个和连接通道维度，然后传递给一个块由3×3深度卷积和批规范化。
通过使用3×3卷积，辅助局部细节可以从q，k，v中聚合。然后利用激活函数和批归一化的线性投影将维数压缩到C，生成细节增强权值。最后，将增强特征与挤压轴向注意给出的特征进行融合。
实验部分将比较不同的包括元素加法和乘法在内的增强模式。3×3深度卷积的时间复杂度为，1×1深度卷积的时间复杂度为。对于激活等其他操作的时间，可以省略。
4、复杂度分析
在应用程序中，设置，以进一步降低计算成本。挤压增强轴向注意力的总时间复杂度为：
如果假设H = W，并取通道为常数。理论上，SEA attention与特征图的大小是线性的。此外，SEA attention只包括卷积、池化、卷积、矩阵乘法等移动友好操作。
5、体系结构和变体
我们介绍了四种变体，SeaFormer Tiny、Small、Base和Large（T、S、B和L）。补充材料中列出了更多配置详细信息。
3、实验
3.1、消融实验
1、SEA attention中各个组件的影响
2、与不同的自注意力模块的比较
3、融合模块宽度的影响
4、模块耗时统计
3.2、SOTA对比
1、ADE20K
2、Cityscapes
3、图像分类
4、参考
[1].SEAFORMER: SQUEEZE-ENHANCED AXIAL TRANSFORMER FOR MOBILE SEMANTIC SEGMENTATION.
5、推荐
[
YOLO涨点Trick | 超越CIOU/SIOU，Wise-IOU让Yolov7再涨1.5个点！](https://mp.weixin.qq.com/s?__biz=MzU5OTA2Mjk5Mw==&mid=2247500890&idx=1&sn=ab2912e05af73cea8354faccd6d1b198&chksm=feb824e4c9cfadf2656924b00ffd2b172fad3ca52cce5faab00487fe6334596bd6b1c9956e00&scene=21#wechat_redirect)
Windows 上基于 TensorRT 的 YOLOV6 部署保姆级教程
扫码加入👉「集智书童-语义分割」交流群
想要了解更多：
前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」
行业技术方案 👉「AI安防、AI医疗、AI自动驾驶」
AI模型部署落地实战 👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」
欢迎扫描上方二维码，加入「 集智书童-知识星球 」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！
点击上方卡片，关注 「集智书童」 公众号