极市导读
在本文中,作者提出了一种移动端友好的架构,名为Token Pyramid Vision Transformer。实验结果表明,TopFormer在多个语义分割数据集上显著优于基于CNN和ViT的网络,并在准确性和实时性之间取得了良好的权衡。>>加入极市CV技术交流群,走在计算机视觉的最前沿
虽然ViT在计算机视觉方面取得了巨大的成功,但巨大的计算成本阻碍了它们在密集的预测任务上的应用,如在移动设备上的语义分割。
在本文中,作者提出了一种移动端友好的架构,名为Token Pyramid Vision Transformer(
TopFormer
)。所提出的最优算法以不同尺度的Token
作为输入,产生尺度感知的语义特征,然后将其注入到相应的Token
中,以增强表征。
实验结果表明,
TopFormer
在多个语义分割数据集上显著优于基于CNN
和ViT
的网络,并在准确性和实时性之间取得了良好的权衡。在ADE20K数据集上,TopFormer
的mIoU比MobileNetV3
的延迟更高5%。此外,TopFormer
的小版本在基于ARM的移动设备上实现实时推理,具有竞争性的结果。
开源地址:https://github.com/hustvl/TopFormer
为了使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所示,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的移动设备上进行实时分割。
其实以及有很多工作对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
,以实现实时的分割任务。
对在移动和嵌入式设备上部署视觉模型的需求不断增加,鼓励了对高效卷积神经网络设计的研究。MobileNet
提出了一种inverted bottleneck
的结构,该结构主要是叠加了Depth-wise
和Point-wise
卷积。IGCNet
和ShuffleNet
使用通道Shuffle/Permutation
操作,为多个组卷积层进行跨Group信息流。GhostNet
使用更简单的操作符,即Depth-wise
卷积,来生成更多的特性。AdderNet
利用add来替换大量的乘法。MobileNeXt
翻转了反向残差块的结构,并呈现了一个连接高维表示的构建块。EfficientNet
和TinyNet
研究了深度、宽度和分辨率的复合尺度。
最精确的分割网络通常需要数十亿个FLOPs的计算,这可能会超过移动设备和嵌入式设备的计算能力。为了加快分割速度和降低计算成本,ICNet
使用多尺度图像作为输入,并使用级联网络来提高计算效率。DFANet
利用一个轻量级的Backbone来加速其网络,并提出了一种跨级特征聚合来提高精度。SwiftNet
使用横向连接作为经济有效的解决方案,在保持速度的同时恢复预测分辨率。BiSeNet
引入了Spatial path
和Semantic path
来减少计算。AlignSeg
和SFNet
对齐了来自相邻level的特征映射,并使用特征金字塔框架进一步增强了特征映射。ESPNets
通过将标准卷积分解为Point-wise convolution
和Spatial pyramid of dilated convolution
来节省计算。
TopFormer
的整体网络架构如图2所示。网络由几个部分组成:
Token Pyramid Module
Semantics Extractor
Semantics Injection Module
Segmentation Head
Token Pyramid Module
将一个图像作为输入,并生成Token Pyramid
。
受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也可以以非常低的计算成本运行。
Scale-aware Semantics Extractor
由几个堆叠的Transformer Block
组成。Transformer Block
数为L。
Transformer Block
由
Multi-head Attention module
、
FFN
和
残差连接
组成。
Token
的空间形状和减少重塑的数量,这里将线性层替换为1×1的卷积层。
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所示,Vision Transformer将来自不同尺度的Token
作为输入。为了进一步减少计算量,使用平均池化操作将不同尺度的Token
数量减少到输入大小的
。来自不同尺度的集合Token
具有相同的分辨率,它们被连接在一起作为Vision Transformer的输入。Vision Transformer可以获得全图像的感受野和丰富的语义。
更具体地说,全局自注意力在空间维度的Token
之间交换信息。1×1卷积层
将在来自不同尺度的Token
之间交换信息。在每个Transformer Block
中,在交换来自所有尺度的Token
信息后学习残差映射,然后将残差映射添加到Token
中,以增强表示和语义。最后,在通过几个Transformer Block
后,获得了尺度感知语义。
在获得尺度感知语义, 直接将它们与其他 Token 相加。然而, 在 Token 和尺 度感知语义之间存在着显著的语义差距。为此, 引入了 Semantics Injection Module 来缓 解在融合这些 Token 之前的语义差距。
如图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个卷积层,生成最终的分割图。
为了定制各种复杂的网络,作者设计了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为目标通道数。各个版本的模型配置如下:
如表所示,将来自不同尺度的堆叠Token
作为语义提取器的输入,并将最后一个Token
分别作为语义提取器的输入。为了公平的比较,附加了一个1×1卷积层来扩展与堆叠的Token
一样的通道。实验结果证明了使用Token Pyramid作为输入的有效性。
如表3所示,使用{1/4,1/8,1/16,1/32}的Token
可以在最重的计算下获得最佳性能。使用{1/16,1/32}的Token
在最轻的计算下获得较差的性能。为了在精度和计算成本之间实现良好的权衡,作者选择在所有其他实验中使用{1/8,1/16,1/32}的Token
。
结果如表所示。在这里使用不带SASE的Topformer作为基线。加入SASE将带来约10%的mIoU收益,这是一个显著的改善。为了验证Transformer Block
中的多头自注意力模块(MHSA),删除了所有的MHSA模块,并添加了更多的ffn,以进行公平的比较。结果表明,在精心的架构设计下是一个高效有效的模块中MHSA可以获得约2.4%的mIoU收益。同时,将SASE与流行的上下文模型进行了比较,如ASPP和PPM。
如表4所示,“+SASE”比“+PSP”和“+ASPP”可以以更低的计算成本获得更好的性能。实验结果表明,SASE更适合用于移动设备。
如表所示,将局部Token
与Sigmoid
层之后的语义相乘,表示为“SigmoidAttn”。将语义提取器中的语义添加到相应的局部Token
中,称为“SemInfo”。与“SigmoidAttn”和“SemInfo”相比,同时添加“SigmoidAttn”和“SemInfo”通过一点额外的计算可以带来很大的改进。
在这里还讨论了Segmentation Head
的设计。将特征传递到Semantic Injection Module
后,输出的层次特征具有较强的语义和丰富的空间细节。提出的Segmentation Head
简单地将它们相加,然后使用2个1×1卷积层
来预测分割图。
作者还设计了另外2个分割头,如图所示。Sum Head
等同于只在SIM中添加SemInfo
。Concat Head
使用1×1卷积层
来减少SIM输出的通道,然后将特征拼接在一起。
如表所示,与Concat head
和Sum head
相比,目前的Segmentation Head
可以取得更好的性能。
如表所示,M=256,192,128通过非常接近的计算实现了类似的性能。因此,在tiny, small和base模型中分别设置M=128,192,256。
不同分辨率的结果如表所示。s32、s64、s128表示集合分辨率为输入大小的 、 。考虑到计算量和精度的权衡性, 选择s64作为语义提取器的输入 Token 的输出 stride。
如图所示,虽然语义提取器具有大部分参数(74%),但语义提取器的FLOPs和实际延迟相对较低(约10%)。
为了进行公平的比较,作者还使用了ImageNet的预训练参数作为初始化。如图所示,提出的TopFormer
的分类架构,将平均池化层和线性层附加到全局语义上,以生成类分数。
由于输入图像的分辨率较小(224×224),这里将语义提取器的输入Token
的目标分辨率设置为输入大小的
。
[1].TopFormer:Token Pyramid Transformer for Mobile Semantic Segmentation
公众号后台回复“数据集”获取30+深度学习数据集下载~
# CV技术社群邀请函 #
备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳)
即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群
每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~