TopFormer：打造Arm端實時分割與檢測模型，完美超越MobileNet!

雖然ViT在計算機視覺方面取得了巨大的成功，但巨大的計算成本阻礙了它們在密集的預(yù)測任務(wù)上的應(yīng)用，如在移動設(shè)備上的語義分割。

在本文中，作者提出了一種移動端友好的架構(gòu)，名為Token Pyramid Vision Transformer(TopFormer)。所提出的最優(yōu)算法以不同尺度的Token作為輸入，產(chǎn)生尺度感知的語義特征，然后將其注入到相應(yīng)的Token中，以增強表征。

實驗結(jié)果表明，TopFormer在多個語義分割數(shù)據(jù)集上顯著優(yōu)于基于CNN和ViT的網(wǎng)絡(luò)，并在準確性和實時性之間取得了良好的權(quán)衡。在ADE20K數(shù)據(jù)集上，TopFormer的mIoU比MobileNetV3的延遲更高5%。此外，TopFormer的小版本在基于ARM的移動設(shè)備上實現(xiàn)實時推理，具有競爭性的結(jié)果。

開源地址：https://github.com/hustvl/TopFormer

1 背景

為了使ViT適應(yīng)各種密集的預(yù)測任務(wù)，最近的ViTs，如PVT、CvT、LeViT以及MobileViT都采用了分層結(jié)構(gòu)，類似的操作也用于卷積神經(jīng)網(wǎng)絡(luò)(CNNs)，如AlexNet和ResNet。這些ViTs將全局自注意力及其變體應(yīng)用到高分辨率Token上，由于Token數(shù)量的二次復雜度，這帶來了巨大的計算成本。

為了提高效率，最近的一些工作，如Swin-Transformer、Shuffle-Transformer、Twins和HR-Former，都在計算在局部/窗口區(qū)域內(nèi)的自注意力。然而，窗口分區(qū)在移動設(shè)備上是非常耗時的。此外，Token slimming和Mobile-Former通過減少Token的數(shù)量而降低了計算能力，但也犧牲了它們的精度。

在這些ViTs中，MobileViT和Mobile-Former是專門為移動設(shè)備設(shè)計的。它們都結(jié)合了CNN和ViT的優(yōu)勢。在圖像分類方面，MobileViT比與參數(shù)數(shù)量相似的MobileNets具有更好的性能。Mobile-Former在比MobileNets更少的FLOPs的情況下取得了更好的性能。然而，與MobileNets相比，它們在移動設(shè)備上的實際延遲方面并沒有顯示出優(yōu)勢。這提出了一個問題：是否有可能設(shè)計出移動友好型網(wǎng)絡(luò)，在移動語義分割任務(wù)上獲得更好的性能？

MobileViT和Mobile-Former的啟發(fā)，作者也利用了CNN和ViT的優(yōu)勢。構(gòu)建了一個基于CNN的模塊，稱為Token Pyramid Module，用于處理高分辨率圖像，以快速生成局部特征金字塔?？紤]到在移動設(shè)備上非常有限的計算能力，在這里使用一些堆疊的輕量級MobileNetV2 Block和Fast Down-Sampling策略來構(gòu)建一個Token Pyramid。

為了獲得豐富的語義和較大的感受野，作者還構(gòu)建了一個基于ViT的模塊，即Semantics Extractor，并將Token作為輸入。為了進一步降低計算成本，使用Average Pooling Operator將Token減少到一個非常小的數(shù)字，例如，輸入大小的1/(64×64)。

與ViT不同，T2T-ViT和LeViT使用嵌入層的最后一個輸出作為輸入Token，而TopFormer將來自不同尺度（階段）的Token池化到非常小的數(shù)字（分辨率），并沿著通道維度進行拼接。然后，新的Token被輸入到Transformer Block中以產(chǎn)生全局語義。由于Transformer Block中的殘差連接學習到的語義與Token的尺度有關(guān)，因此該模塊被表示為Scale-aware Global Semantics。

為了獲得密集預(yù)測任務(wù)的強大層次特征，將尺度感知的全局語義通過不同尺度的Token通道進行分割，然后將標度感知的全局語義與相應(yīng)的Token融合，以增強表示。增強的Token被用作分割Head的輸入。

為了證明方法的有效性，在具有挑戰(zhàn)性的分割數(shù)據(jù)集上進行了實驗：ADE20K，Pascal上下文和COCOStuff。并測試了硬件上的延遲，即一個現(xiàn)成的基于Arm的計算核心。

如圖1所示，TopFormer比延遲較低的MobileNets獲得了更好的結(jié)果。為了證明方法的泛化性,還在COCO數(shù)據(jù)集上進行了目標檢測實驗。

綜上所述，本文的貢獻如下：

1. 所提出的最優(yōu)預(yù)測算法以不同尺度的Token作為輸入，并將Token池化到非常小的尺寸，以獲得計算代價非常輕的Scale-aware Global Semantics;
2. 所提出的Semantics Injection Module可以將Scale-aware Global Semantics注入到相應(yīng)的Token中，構(gòu)建強大的層次特征;
3. 與MobileNetV3相比，TopFormer可以實現(xiàn)5%的mIoU提升，在ADE20K數(shù)據(jù)集上基于Arm的移動設(shè)備上的延遲更低。TopFormer-Tiny可以在基于Arm的移動設(shè)備上進行實時分割。

2 相關(guān)工作

2.1 輕量化Vision Transformers

其實以及有很多工作對Vision Transformers結(jié)構(gòu)在圖像識別中的應(yīng)用進行了探索。ViT是第一個將純Transformer應(yīng)用于圖像分類的工作，實現(xiàn)了最先進的性能。隨后，DeiT引入了基于Token的蒸餾，以減少訓練Transformer所需的數(shù)據(jù)量。T2T-ViT通過遞歸地將相鄰Token聚合為一個Token來減少Token長度。Swin-Transformer在每個局部窗口內(nèi)計算自注意力，帶來了輸入Token數(shù)量的線性計算復雜度。然而，這些Vision Transformers和后續(xù)工作往往是大量的參數(shù)和沉重的計算復雜度。

為了構(gòu)建一個輕量級的Vision Transformer，LeViT設(shè)計了一個混合架構(gòu)，使用堆疊的標準卷積層和stride=2來減少Token的數(shù)量，然后附加一個改進的Vision Transformer來提取語義。在分類任務(wù)中，LeViT在CPU上的性能明顯優(yōu)于EfficientNet。

MobileViT也采用了相同的策略，并使用MobilenetV2 Block而不是標準的卷積層來對特征圖進行降采樣。Mobile-Former采用并行結(jié)構(gòu)與雙向橋，利用了MobileNet和Transformer的優(yōu)勢。然而，有研究表明MobileViT和其他ViT-based的網(wǎng)絡(luò)在移動設(shè)備上明顯慢于MobileNets。

對于分割任務(wù)，輸入的圖像總是高分辨率的。因此，ViT-based網(wǎng)絡(luò)比MobileNets的執(zhí)行速度更具挑戰(zhàn)性。本文中的目標是設(shè)計一個輕量級的視Vision Transformer，同時改模型可以超過MobileNets，以實現(xiàn)實時的分割任務(wù)。

2.2 高效的卷積神經(jīng)網(wǎng)絡(luò)

對在移動和嵌入式設(shè)備上部署視覺模型的需求不斷增加，鼓勵了對高效卷積神經(jīng)網(wǎng)絡(luò)設(shè)計的研究。MobileNet提出了一種inverted bottleneck的結(jié)構(gòu)，該結(jié)構(gòu)主要是疊加了Depth-wise和Point-wise卷積。IGCNet和ShuffleNet使用通道Shuffle/Permutation操作，為多個組卷積層進行跨Group信息流。GhostNet使用更簡單的操作符，即Depth-wise卷積，來生成更多的特性。AdderNet利用add來替換大量的乘法。MobileNeXt翻轉(zhuǎn)了反向殘差塊的結(jié)構(gòu)，并呈現(xiàn)了一個連接高維表示的構(gòu)建塊。EfficientNet和TinyNet研究了深度、寬度和分辨率的復合尺度。

2.3 移動端語義分割

最精確的分割網(wǎng)絡(luò)通常需要數(shù)十億個FLOPs的計算，這可能會超過移動設(shè)備和嵌入式設(shè)備的計算能力。為了加快分割速度和降低計算成本，ICNet使用多尺度圖像作為輸入，并使用級聯(lián)網(wǎng)絡(luò)來提高計算效率。DFANet利用一個輕量級的Backbone來加速其網(wǎng)絡(luò)，并提出了一種跨級特征聚合來提高精度。SwiftNet使用橫向連接作為經(jīng)濟有效的解決方案，在保持速度的同時恢復預(yù)測分辨率。BiSeNet引入了Spatial path和Semantic path來減少計算。AlignSeg和SFNet對齊了來自相鄰level的特征映射，并使用特征金字塔框架進一步增強了特征映射。ESPNets通過將標準卷積分解為Point-wise convolution和Spatial pyramid of dilated convolution來節(jié)省計算。

3 本文方法

TopFormer的整體網(wǎng)絡(luò)架構(gòu)如圖2所示。網(wǎng)絡(luò)由幾個部分組成：

• Token Pyramid Module
• Semantics Extractor
• Semantics Injection Module
• Segmentation Head

Token Pyramid Module將一個圖像作為輸入，并生成Token Pyramid。

3.1 Token Pyramid Module

受MobileNets的啟發(fā)，所提出的Token Pyramid Module由堆疊的MobileNet blocks組成。與MobileNets不同，Token Pyramid Module的目標并不是獲得豐富的語義和較大的感受野，而是使用更少的塊來構(gòu)建Token Pyramid。

如圖2所示，把一個圖像作為輸入，其中3，H，W分別表示RGB通道，高度，寬度；

Token Pyramid Module:

• 首先，通過一些MobileNetV2 Block產(chǎn)生一系列Token，N表示Scale的數(shù)量。
• 然后，將Token平均池化到目標大小，例如，。
• 最后，將來自不同尺度的Token沿著通道維度連接起來，產(chǎn)生新的Token。新的Token將被輸入Vision Transformer，以產(chǎn)生具有尺度感知的語義特征。

由于新的Token的數(shù)量較小，因此即使新的Token具有較大的通道，Vision Transformer也可以以非常低的計算成本運行。

3.2 Scale-aware Semantics Extractor

Scale-aware Semantics Extractor由幾個堆疊的Transformer Block組成。Transformer Block數(shù)為L。

• Transformer Block由Multi-head Attention module、FFN和殘差連接組成。
• 為了保持Token的空間形狀和減少重塑的數(shù)量，這里將線性層替換為1×1的卷積層。
• 此外，在ViT中，所有的非線性激活都是ReLU6，而不是GELU。

對于Multi-head Attention module，遵循LeViT的配置，將key K和query Q的Head尺寸設(shè)置為D=16，value V的head 設(shè)置為2D=32通道。在計算Attention Map和輸出時，減少K和Q的通道將降低計算成本。同時，還去掉了Layer Normalization Layer，并向每個卷積添加了Batch Normalization。在推理過程中，Batch Normalization可以與前面的卷積融合。

對于FFN，通過在2個1×1卷積層之間插入一個Depth-wise卷積，來增強Vision Transformer的局部連接。將FFN的擴展系數(shù)設(shè)為2，以降低計算成本。Transformer Block的數(shù)量是L。

如圖3所示，Vision Transformer將來自不同尺度的Token作為輸入。為了進一步減少計算量，使用平均池化操作將不同尺度的Token數(shù)量減少到輸入大小的。來自不同尺度的集合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(SIM)以Token Pyramid Module的局部Token和Vision Transformer的全局語義作為輸入。

局部Token通過1×1卷積層，然后進行批歸一化，生成要注入的特征。

全局語義輸入1×1卷積層 + 批歸一化層 + sigmoid層產(chǎn)生語義權(quán)重，同時全局語義也通過1×1卷積層 + 批歸一化。

這3個輸出的大小相同。然后，通過阿達瑪生產(chǎn)將全局語義注入到局部標記中，并在注入后添加全局語義和特征。幾個sim的輸出共享相同的通道數(shù)，記為M。經(jīng)過Semantics Injection Module后，來自不同尺度的增強Token同時捕獲了豐富的空間信息和語義信息，這對語義分割至關(guān)重要。此外，Semantics Injection Module還緩解了Token之間的語義差距。所提出的Segmentation Head首先將低分辨率Token上采樣到與高分辨率Token相同的大小，并按元素方式對所有尺度的Token進行sum up。最后，將該特征通過2個卷積層，生成最終的分割圖。

3.4 架構(gòu)及其變體

為了定制各種復雜的網(wǎng)絡(luò)，作者設(shè)計了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為目標通道數(shù)。各個版本的模型配置如下：

4 實驗

4.1 消融實驗

1、Token Pyramid的影響

如表所示，將來自不同尺度的堆疊Token作為語義提取器的輸入，并將最后一個Token分別作為語義提取器的輸入。為了公平的比較，附加了一個1×1卷積層來擴展與堆疊的Token一樣的通道。實驗結(jié)果證明了使用Token Pyramid作為輸入的有效性。

如表3所示，使用{1/4,1/8,1/16,1/32}的Token可以在最重的計算下獲得最佳性能。使用{1/16,1/32}的Token在最輕的計算下獲得較差的性能。為了在精度和計算成本之間實現(xiàn)良好的權(quán)衡，作者選擇在所有其他實驗中使用{1/8,1/16,1/32}的Token。

2、Scale-aware Semantics Extractor的影響

結(jié)果如表所示。在這里使用不帶SASE的Topformer作為基線。加入SASE將帶來約10%的mIoU收益，這是一個顯著的改善。為了驗證Transformer Block中的多頭自注意力模塊(MHSA)，刪除了所有的MHSA模塊，并添加了更多的ffn，以進行公平的比較。結(jié)果表明，在精心的架構(gòu)設(shè)計下是一個高效有效的模塊中MHSA可以獲得約2.4%的mIoU收益。同時，將SASE與流行的上下文模型進行了比較，如ASPP和PPM。

如表4所示，“+SASE”比“+PSP”和“+ASPP”可以以更低的計算成本獲得更好的性能。實驗結(jié)果表明，SASE更適合用于移動設(shè)備。

3、Semantic Injection Module和Segmentation Head的影響

如表所示，將局部Token與Sigmoid層之后的語義相乘，表示為“SigmoidAttn”。將語義提取器中的語義添加到相應(yīng)的局部Token中，稱為“SemInfo”。與“SigmoidAttn”和“SemInfo”相比，同時添加“SigmoidAttn”和“SemInfo”通過一點額外的計算可以帶來很大的改進。

在這里還討論了Segmentation Head的設(shè)計。將特征傳遞到Semantic Injection Module后，輸出的層次特征具有較強的語義和豐富的空間細節(jié)。提出的Segmentation Head簡單地將它們相加，然后使用2個1×1卷積層來預(yù)測分割圖。

作者還設(shè)計了另外2個分割頭，如圖所示。Sum Head等同于只在SIM中添加SemInfo。Concat Head使用1×1卷積層來減少SIM輸出的通道，然后將特征拼接在一起。

如表所示，與Concat head和Sum head相比，目前的Segmentation Head可以取得更好的性能。

4、SIM寬度的影響

如表所示，M=256,192,128通過非常接近的計算實現(xiàn)了類似的性能。因此，在tiny, small和base模型中分別設(shè)置M=128,192,256。

5、output stride的影響

不同分辨率的結(jié)果如表所示。s32、s64、s128表示集合分辨率為輸入大小的 、 。考慮到計算量和精度的權(quán)衡性, 選擇s64作為語義提取器的輸入 Token 的輸出 stride。

6、參數(shù)量與實時性

如圖所示，雖然語義提取器具有大部分參數(shù)（74%），但語義提取器的FLOPs和實際延遲相對較低（約10%）。

4.2 圖像分類

為了進行公平的比較，作者還使用了ImageNet的預(yù)訓練參數(shù)作為初始化。如圖所示，提出的TopFormer的分類架構(gòu)，將平均池化層和線性層附加到全局語義上，以生成類分數(shù)。

由于輸入圖像的分辨率較?。?24×224），這里將語義提取器的輸入Token的目標分辨率設(shè)置為輸入大小的。

4.3 語義分割

1、ADE20K

2、Cityscapes

3、可視化結(jié)果

4.4 目標檢測

參考

[1].TopFormer：Token Pyramid Transformer for Mobile Semantic Segmentation

公眾號后臺回復“數(shù)據(jù)集”獲取30+深度學習數(shù)據(jù)集下載～