1背景
为了使ViT适应各种密集的预测任务,最近的ViTs,如PVT
、CvT
、LeViT
以及MobileViT
都采用了分层结构,类似的操作也用于卷积神经网络(CNNs),如AlexNet
和ResNet
。这些ViTs将全局自注意力及其变体应用到高分辨率Token
上,由于Token
数量的二次复杂度,这带来了巨大的计算成本。
为了提高效率,最近的一些工作,如Swin-Transformer
、Shuffle-Transformer
、Twins
和HR-Former
,都在计算在局部/窗口区域内的自注意力。然而,窗口分区在移动设备上是非常耗时的。此外,Token slimming
和Mobile-Former
通过减少Token
的数量而降低了计算能力,但也牺牲了它们的精度。
在这些ViTs中,MobileViT
和Mobile-Former
是专门为移动设备设计的。它们都结合了CNN和ViT的优势。在图像分类方面,MobileViT
比与参数数量相似的MobileNets
具有更好的性能。Mobile-Former
在比MobileNets
更少的FLOPs的情况下取得了更好的性能。然而,与MobileNets
相比,它们在移动设备上的实际延迟方面并没有显示出优势。这提出了一个问题:是否有可能设计出移动友好型网络,在移动语义分割任务上获得更好的性能?
MobileViT
和Mobile-Former
的启发,作者也利用了CNN和ViT的优势。构建了一个基于CNN的模块,称为Token Pyramid Module
,用于处理高分辨率图像,以快速生成局部特征金字塔。考虑到在移动设备上非常有限的计算能力,在这里使用一些堆叠的轻量级MobileNetV2 Block
和Fast Down-Sampling
策略来构建一个Token Pyramid
。
为了获得丰富的语义和较大的感受野,作者还构建了一个基于ViT的模块,即Semantics Extractor
,并将Token
作为输入。为了进一步降低计算成本,使用Average Pooling Operator
将Token
减少到一个非常小的数字,例如,输入大小的1/(64×64)。
与ViT
不同,T2T-ViT
和LeViT
使用嵌入层的最后一个输出作为输入Token
,而TopFormer
将来自不同尺度(阶段)的Token
池化到非常小的数字(分辨率),并沿着通道维度进行拼接。然后,新的Token
被输入到Transformer Block
中以产生全局语义。由于Transformer Block
中的残差连接学习到的语义与Token
的尺度有关,因此该模块被表示为Scale-aware Global Semantics
。
为了获得密集预测任务的强大层次特征,将尺度感知的全局语义通过不同尺度的Token
通道进行分割,然后将标度感知的全局语义与相应的Token
融合,以增强表示。增强的Token
被用作分割Head的输入。
为了证明方法的有效性,在具有挑战性的分割数据集上进行了实验:ADE20K,Pascal上下文和COCOStuff。并测试了硬件上的延迟,即一个现成的基于Arm的计算核心。
图1
如图1所示,TopFormer
比延迟较低的MobileNets
获得了更好的结果。为了证明方法的泛化性,还在COCO数据集上进行了目标检测实验。
综上所述,本文的贡献如下:
- 所提出的最优预测算法以不同尺度的
Token
作为输入,并将Token
池化到非常小的尺寸,以获得计算代价非常轻的Scale-aware Global Semantics
; - 所提出的
Semantics Injection Module
可以将Scale-aware Global Semantics
注入到相应的Token
中,构建强大的层次特征; - 与
MobileNetV3
相比,TopFormer
可以实现5%的mIoU提升,在ADE20K数据集上基于Arm的移动设备上的延迟更低。TopFormer-Tiny
可以在基于Arm的移动设备上进行实时分割。
2相关工作
2.1 轻量化Vision Transformers
其实以及有很多工作对Vision Transformers结构在图像识别中的应用进行了探索。ViT
是第一个将纯Transformer应用于图像分类的工作,实现了最先进的性能。随后,DeiT
引入了基于Token
的蒸馏,以减少训练Transformer所需的数据量。T2T-ViT
通过递归地将相邻Token
聚合为一个Token
来减少Token
长度。Swin-Transformer
在每个局部窗口内计算自注意力,带来了输入Token
数量的线性计算复杂度。然而,这些Vision Transformers和后续工作往往是大量的参数和沉重的计算复杂度。
为了构建一个轻量级的Vision Transformer,LeViT
设计了一个混合架构,使用堆叠的标准卷积层和stride=2来减少Token
的数量,然后附加一个改进的Vision Transformer来提取语义。在分类任务中,LeViT
在CPU上的性能明显优于EfficientNet
。
MobileViT
也采用了相同的策略,并使用MobilenetV2 Block
而不是标准的卷积层来对特征图进行降采样。Mobile-Former
采用并行结构与双向桥,利用了MobileNet
和Transformer
的优势。然而,有研究表明MobileViT
和其他ViT-based
的网络在移动设备上明显慢于MobileNets
。
对于分割任务,输入的图像总是高分辨率的。因此,ViT-based
网络比MobileNets
的执行速度更具挑战性。本文中的目标是设计一个轻量级的视Vision Transformer,同时改模型可以超过MobileNets
,以实现实时的分割任务。
2.2 高效的卷积神经网络
对在移动和嵌入式设备上部署视觉模型的需求不断增加,鼓励了对高效卷积神经网络设计的研究。MobileNet
提出了一种inverted bottleneck
的结构,该结构主要是叠加了Depth-wise
和Point-wise
卷积。IGCNet
和ShuffleNet
使用通道Shuffle/Permutation
操作,为多个组卷积层进行跨Group信息流。GhostNet
使用更简单的操作符,即Depth-wise
卷积,来生成更多的特性。AdderNet
利用add来替换大量的乘法。MobileNeXt
翻转了反向残差块的结构,并呈现了一个连接高维表示的构建块。EfficientNet
和TinyNet
研究了深度、宽度和分辨率的复合尺度。
2.3 移动端语义分割
最精确的分割网络通常需要数十亿个FLOPs的计算,这可能会超过移动设备和嵌入式设备的计算能力。为了加快分割速度和降低计算成本,ICNet
使用多尺度图像作为输入,并使用级联网络来提高计算效率。DFANet
利用一个轻量级的Backbone来加速其网络,并提出了一种跨级特征聚合来提高精度。SwiftNet
使用横向连接作为经济有效的解决方案,在保持速度的同时恢复预测分辨率。BiSeNet
引入了Spatial path
和Semantic path
来减少计算。AlignSeg
和SFNet
对齐了来自相邻level的特征映射,并使用特征金字塔框架进一步增强了特征映射。ESPNets
通过将标准卷积分解为Point-wise convolution
和Spatial pyramid of dilated convolution
来节省计算。
3本文方法
TopFormer
的整体网络架构如图2所示。网络由几个部分组成:
Token Pyramid Module
Semantics Extractor
Semantics Injection Module
Segmentation Head
Token Pyramid Module
将一个图像作为输入,并生成Token Pyramid
。
图2.TopFormer
的整体架构
3.1 Token Pyramid Module
受MobileNets
的启发,所提出的Token Pyramid Module
由堆叠的MobileNet blocks
组成。与MobileNets
不同,Token Pyramid Module
的目标并不是获得丰富的语义和较大的感受野,而是使用更少的块来构建Token Pyramid
。
如图2所示,把一个图像作为输入,其中3,H,W分别表示RGB通道,高度,宽度;
Token Pyramid Module
:
- 首先,通过一些
MobileNetV2 Block
产生一系列Token
,N表示Scale的数量。 - 然后,将
Token
平均池化到目标大小,例如,。 - 最后,将来自不同尺度的
Token
沿着通道维度连接起来,产生新的Token
。新的Token
将被输入Vision Transformer,以产生具有尺度感知的语义特征。
由于新的Token
的数量较小,因此即使新的Token
具有较大的通道,Vision Transformer也可以以非常低的计算成本运行。
3.2 Scale-aware Semantics Extractor
Scale-aware Semantics Extractor
由几个堆叠的Transformer Block
组成。Transformer Block
数为L。
Transformer Block
由Multi-head Attention module
、FFN
和残差连接
组成。- 为了保持
Token
的空间形状和减少重塑的数量,这里将线性层替换为1×1的卷积层。 - 此外,在ViT中,所有的非线性激活都是
ReLU6
,而不是GELU
。
对于Multi-head Attention module
,遵循LeViT的配置,将key K和query Q的Head尺寸设置为D=16,value V的head 设置为2D=32通道。在计算Attention Map和输出时,减少K和Q的通道将降低计算成本。同时,还去掉了Layer Normalization Layer
,并向每个卷积添加了Batch Normalization
。在推理过程中,Batch Normalization
可以与前面的卷积融合。
对于FFN
,通过在2个1×1卷积
层之间插入一个Depth-wise
卷积,来增强Vision Transformer的局部连接。将FFN的扩展系数设为2,以降低计算成本。Transformer Block
的数量是L。
图3.Scale-aware Semantics Extractor
如图3所示,Vision Transformer将来自不同尺度的Token
作为输入。为了进一步减少计算量,使用平均池化操作将不同尺度的Token
数量减少到输入大小的。来自不同尺度的集合Token
具有相同的分辨率,它们被连接在一起作为Vision Transformer的输入。Vision Transformer可以获得全图像的感受野和丰富的语义。
更具体地说,全局自注意力在空间维度的Token
之间交换信息。1×1卷积层
将在来自不同尺度的Token
之间交换信息。在每个Transformer Block
中,在交换来自所有尺度的Token
信息后学习残差映射,然后将残差映射添加到Token
中,以增强表示和语义。最后,在通过几个Transformer Block
后,获得了尺度感知语义。
3.3 Semantics Injection Module and Segmentation Head
在获得尺度感知语义,直接将它们与其他Token
相加。然而,在Token
和尺度感知语义之间存在着显著的语义差距。为此,引入了Semantics Injection Module
来缓解在融合这些Token
之前的语义差距。
图4.Semantics Injection Module和Segmentation Head
如图4所示,
Semantics Injection Module
(SIM)以Token Pyramid Module
的局部Token
和Vision Transformer的全局语义作为输入。局部
Token
通过1×1卷积层
,然后进行批归一化
,生成要注入的特征。全局语义输入
1×1卷积层
+批归一化层
+sigmoid层
产生语义权重,同时全局语义也通过1×1卷积层
+批归一化
。
这3个输出的大小相同。然后,通过阿达玛生产将全局语义注入到局部标记中,并在注入后添加全局语义和特征。几个sim的输出共享相同的通道数,记为M。
经过Semantics Injection Module
后,来自不同尺度的增强Token
同时捕获了丰富的空间信息和语义信息,这对语义分割至关重要。此外,Semantics Injection Module
还缓解了Token
之间的语义差距。所提出的Segmentation Head
首先将低分辨率Token
上采样到与高分辨率Token
相同的大小,并按元素方式对所有尺度的Token
进行sum up
。最后,将该特征通过2个卷积层,生成最终的分割图。
3.4 架构及其变体
为了定制各种复杂的网络,作者设计了TopFormer-Tiny
(TopFormer-T)和TopFormer-Small
(TopFormer-S)和TopFormer-Base
(TopFormer-Base)。
下表给出了Base、Small以及Tiny模型的尺寸和FLOPs。Base、Small以及Tiny模型在每个Multi-Head self-attention module
中分别有8、6和4个Head,以M=256、M=192和M=128为目标通道数。各个版本的模型配置如下:
模型配置表