Transformer崛起| TopFormer打造Arm端實(shí)時(shí)分割與檢測(cè)模型,完美超...

雖然ViT在計(jì)算機(jī)視覺方面取得了巨大的成功，但巨大的計(jì)算成本阻礙了它們?cè)诿芗念A(yù)測(cè)任務(wù)上的應(yīng)用，如在移動(dòng)設(shè)備上的語(yǔ)義分割。

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

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

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

1背景

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

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

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

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

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

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

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

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

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

綜上所述，本文的貢獻(xiàn)如下：

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

2相關(guān)工作

2.1 輕量化Vision Transformers

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

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

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

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

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

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

2.3 移動(dòng)端語(yǔ)義分割

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

3本文方法

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

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

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

3.1 Token Pyramid Module

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

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

Token Pyramid Module:

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

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

3.2 Scale-aware Semantics Extractor

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

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

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

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

如圖3所示，Vision Transformer將來自不同尺度的Token作為輸入。為了進(jìn)一步減少計(jì)算量，使用平均池化操作將不同尺度的Token數(shù)量減少到輸入大小的。來自不同尺度的集合Token具有相同的分辨率，它們被連接在一起作為Vision Transformer的輸入。Vision Transformer可以獲得全圖像的感受野和豐富的語(yǔ)義。

更具體地說，全局自注意力在空間維度的Token之間交換信息。1×1卷積層將在來自不同尺度的Token之間交換信息。在每個(gè)Transformer Block中，在交換來自所有尺度的Token信息后學(xué)習(xí)殘差映射，然后將殘差映射添加到Token中，以增強(qiáng)表示和語(yǔ)義。最后，在通過幾個(gè)Transformer Block后，獲得了尺度感知語(yǔ)義。

3.3 Semantics Injection Module and Segmentation Head

在獲得尺度感知語(yǔ)義，直接將它們與其他Token 相加。然而，在Token和尺度感知語(yǔ)義之間存在著顯著的語(yǔ)義差距。為此，引入了Semantics Injection Module來緩解在融合這些Token之前的語(yǔ)義差距。

如圖4所示，Semantics Injection Module(SIM)以Token Pyramid Module的局部Token和Vision Transformer的全局語(yǔ)義作為輸入。

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

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

這3個(gè)輸出的大小相同。然后，通過阿達(dá)瑪生產(chǎn)將全局語(yǔ)義注入到局部標(biāo)記中，并在注入后添加全局語(yǔ)義和特征。幾個(gè)sim的輸出共享相同的通道數(shù)，記為M。

經(jīng)過Semantics Injection Module后，來自不同尺度的增強(qiáng)Token同時(shí)捕獲了豐富的空間信息和語(yǔ)義信息，這對(duì)語(yǔ)義分割至關(guān)重要。此外，Semantics Injection Module還緩解了Token之間的語(yǔ)義差距。所提出的Segmentation Head首先將低分辨率Token上采樣到與高分辨率Token相同的大小，并按元素方式對(duì)所有尺度的Token進(jìn)行sum up。最后，將該特征通過2個(gè)卷積層，生成最終的分割圖。

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

為了定制各種復(fù)雜的網(wǎng)絡(luò)，作者設(shè)計(jì)了TopFormer-Tiny (TopFormer-T)和TopFormer-Small(TopFormer-S)和TopFormer-Base(TopFormer-Base)。

下表給出了Base、Small以及Tiny模型的尺寸和FLOPs。Base、Small以及Tiny模型在每個(gè)Multi-Head self-attention module中分別有8、6和4個(gè)Head，以M=256、M=192和M=128為目標(biāo)通道數(shù)。各個(gè)版本的模型配置如下：

4實(shí)驗(yàn)

4.1 消融實(shí)驗(yàn)

1、Token Pyramid的影響

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

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

2、Scale-aware Semantics Extractor的影響

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

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

3、Semantic Injection Module和Segmentation Head的影響

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

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

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

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

4、SIM寬度的影響

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

5、output stride的影響

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

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

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

4.2 圖像分類

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

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

4.3 語(yǔ)義分割

1、ADE20K

2、Cityscapes

3、可視化結(jié)果

4.4 目標(biāo)檢測(cè)

5參考

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

6推薦閱讀

阿里巴巴提出USI 讓AI煉丹自動(dòng)化了,訓(xùn)練任何Backbone無(wú)需超參配置,實(shí)現(xiàn)大一統(tǒng)！

超快語(yǔ)義分割 | PP-LiteSeg集速度快、精度高、易部署等優(yōu)點(diǎn)于一身，必會(huì)模型！??！

CVPR2022 oral | MetaFormer才是探索Transformer的源泉,衍生PoolFormer速度喜人

長(zhǎng)按掃描下方二維碼添加小助手。

可以一起討論遇到的問題

聲明：轉(zhuǎn)載請(qǐng)說明出處

掃描下方二維碼關(guān)注【集智書童】公眾號(hào)，獲取更多實(shí)踐項(xiàng)目源碼和論文解讀，非常期待你我的相遇，讓我們以夢(mèng)為馬，砥礪前行！