TPAMI 2022｜金字塔池化的骨干網(wǎng)絡(luò)，各大任務(wù)都漲點(diǎn)！南開(kāi)&達(dá)摩院聯(lián)合推出P2T

作者丨吳宇寰

編輯丨極市平臺(tái)

導(dǎo)讀

 

該工作主要受PVT和MViT的啟發(fā)，第一次將金字塔池化融入到視覺(jué)transformer的骨干網(wǎng)絡(luò)中，在減小輸入序列長(zhǎng)度的同時(shí)去捕捉多層次的豐富特征表達(dá)。P2T在圖像分類、語(yǔ)義分割、目標(biāo)檢測(cè)、實(shí)例分割等多個(gè)領(lǐng)域中，并取得了比現(xiàn)有CNN/Transformer骨干更優(yōu)異的性能。

作者列表：吳宇寰，劉云，占新，程明明

作者單位：南開(kāi)大學(xué)，阿里達(dá)摩院

代碼位置：https://github.com/yuhuan-wu/P2T

論文地址：https://arxiv.org/abs/2106.12011

IEEE地址：https://ieeexplore.ieee.org/document/9870559

中文版本：https://mmcheng.net/wp-content/uploads/2022/08/22PAMI_P2T_CN.pdf

1. 引言Transformer技術(shù)在計(jì)算機(jī)視覺(jué)領(lǐng)域的興起

近年來(lái)，NLP中的transformer技術(shù)在計(jì)算機(jī)視覺(jué)領(lǐng)域掀起新的熱潮。Transformer可以充分地捕捉長(zhǎng)距離的全局特征關(guān)系，正好滿足了計(jì)算機(jī)視覺(jué)大感受野的需求。視覺(jué)transformer [1] 也在近年來(lái)的各大領(lǐng)域取得了巨大的成功，證明了純transformer架構(gòu)在巨大的數(shù)據(jù)規(guī)模下訓(xùn)練也能在大規(guī)模的ImageNet分類數(shù)據(jù)集上與CNN方法取得相同的分類效果。DeiT [2] 則隨后證明了使用知識(shí)蒸餾可以只在ImgeNet-1K分類數(shù)據(jù)集上就能與CNN取得類似的效果。

然而與NLP領(lǐng)域特征長(zhǎng)度有限不同，計(jì)算機(jī)視覺(jué)的圖像具有密集特性，其拉直(flattened)后的特征長(zhǎng)度要遠(yuǎn)遠(yuǎn)大于NLP任務(wù)中的序列長(zhǎng)度。比如如WMT2014數(shù)據(jù)集 [3] 的平均序列長(zhǎng)度為25，ImageNet-1K數(shù)據(jù)集 [4] 的一般使用圖像分辨率為224x224的大小，即使使用常規(guī)的倍降采樣，拉直后的圖像序列長(zhǎng)度仍有3136。而Transformer因其需要捕捉充分的全局注意力關(guān)系，其多頭自我注意力（multi-head self-attention, MHSA）的計(jì)算復(fù)雜度是與輸入特征呈平方關(guān)系。CNN網(wǎng)絡(luò)因?yàn)橹饕蹲骄植啃畔?，其?jì)算復(fù)雜度只與圖像的特征長(zhǎng)度呈線性相關(guān)。

因此，在計(jì)算機(jī)視覺(jué)領(lǐng)域中應(yīng)用transformer技術(shù)的關(guān)鍵困難在于如何使計(jì)算注意力關(guān)系的過(guò)程更加高效。視覺(jué)transformer [1]作為第一個(gè)解決方案，首先提出使用大的降采樣尺度（如16、32倍）來(lái)對(duì)輸入特征進(jìn)行降采樣，并設(shè)計(jì)了一個(gè)純transformer的網(wǎng)絡(luò)骨干（backbone），并在圖像分類任務(wù)上取得了巨大的成功。而隨后PVT [5]、Swin [6]、MViT [7] 等金字塔結(jié)構(gòu)的網(wǎng)絡(luò)骨干使用類似CNN的逐漸下采樣的手段，來(lái)支持各類下游視覺(jué)任務(wù)。

目前transformer網(wǎng)絡(luò)骨干的思路

如上文所述，transformer應(yīng)用到計(jì)算機(jī)視覺(jué)領(lǐng)域的關(guān)鍵問(wèn)題在于如何使計(jì)算MHSA的過(guò)程更加高效。PVT和MViT對(duì)特征圖采用了單個(gè)池化操作來(lái)減少M(fèi)HSA的計(jì)算量和顯存占用，相應(yīng)地是進(jìn)行輸入特征圖元素與不同區(qū)域之間的關(guān)系建模，而不是特征圖元素之間的關(guān)系建模。Swin transformer則是另一種思路：它先將特征劃分為小的區(qū)塊，并在每個(gè)小區(qū)塊內(nèi)直接計(jì)算MHSA，即進(jìn)行小區(qū)塊內(nèi)元素之間的關(guān)系建模；隨后，它將區(qū)塊劃分進(jìn)行偏移，同時(shí)計(jì)算偏移后的小區(qū)塊的MHSA。這樣，類似CNN網(wǎng)絡(luò)一樣，隨著網(wǎng)絡(luò)的堆疊，Swin transformer網(wǎng)絡(luò)的感受野會(huì)逐步增強(qiáng)。然而，視覺(jué)transformer的最根本特性源于其對(duì)特征全局的關(guān)系建模，這也是為什么我們要嘗試將transformer引入到過(guò)度由CNN統(tǒng)治的計(jì)算機(jī)視覺(jué)領(lǐng)域中。

本文的工作

我們的工作主要受PVT和MViT的啟發(fā)，即使用特征降采樣的形式來(lái)減少M(fèi)HSA的計(jì)算量。PVT和MViT使用單個(gè)尺度的降采樣后的特征來(lái)計(jì)算MHSA，這種特征表達(dá)的模式有限。我們想到，如果使用金字塔池化（pyramid pooling）來(lái)生成特征，既能夠大大縮小特征序列長(zhǎng)度，又能引入更豐富的多層次特征。

金字塔池化在計(jì)算機(jī)視覺(jué)中有著悠久的歷史 [8-9]。它利用多個(gè)不同步長(zhǎng)的池化操作來(lái)提取豐富不同感受野的池化特征，其在目標(biāo)檢測(cè) [10] 和語(yǔ)義分割中 [11] 也被證明十分有效。然而，它們都依賴于一個(gè)CNN骨干網(wǎng)絡(luò)，所以它們也只是被專門設(shè)計(jì)用于某個(gè)特定的任務(wù)。目前并沒(méi)有探索金字塔池化在骨干網(wǎng)絡(luò)設(shè)計(jì)的工作。而一個(gè)新型的骨干網(wǎng)絡(luò)可以應(yīng)用于多種下游任務(wù)中。我們第一次將金字塔池化融入到視覺(jué)transformer的骨干網(wǎng)絡(luò)中，在減小輸入序列長(zhǎng)度的同時(shí)去捕捉多層次的豐富特征表達(dá)。

利用金字塔池化，我們?cè)O(shè)計(jì)了基于的金字塔池化的MHSA（pooling-based MHSA, P-MHSA）。金字塔池化作為計(jì)算P-MHSA中降低序列長(zhǎng)度的基本操作。使用P-MHSA替換掉MHSA后，即為P2T基礎(chǔ)模塊。通過(guò)堆疊不同數(shù)量的P2T基礎(chǔ)模塊，我們?cè)O(shè)計(jì)了P2T-Tiny/Small/Base/Large骨干網(wǎng)絡(luò)，分別對(duì)應(yīng)ResNet-18/50/101/152。P2T系列的多層次特征捕捉能力也要強(qiáng)于之前的PVT與MViT方法。我們也在圖像分類、語(yǔ)義分割、目標(biāo)檢測(cè)、實(shí)例分割等多個(gè)領(lǐng)域中，并取得了比現(xiàn)有CNN/Transformer骨干更優(yōu)異的性能，如ResNet, ResNeXt, Res2Net, PVT, Swin, Twins, 及PVTv2。

圖 1 P2T與其他方法在ADE20K數(shù)據(jù)集上的對(duì)比圖（參數(shù)量、mIoU精度）。

2. P2T 方法

P2T的總體結(jié)構(gòu)如圖 2所示。每個(gè)階段開(kāi)始時(shí)有一個(gè)patch embedding模塊，用于對(duì)輸入的特征進(jìn)行下采樣。第一個(gè)Patch Embedding模塊通過(guò)步長(zhǎng)為4的7x7 卷積將輸入特征下采樣為原來(lái)的1/4，其他Patch Embedding則是將輸入特征下采樣為原來(lái)的2倍。每個(gè)Patch Embedding模塊之后則堆疊了不同數(shù)量的P2T基礎(chǔ)模塊。其中，P2T基本模塊的計(jì)算可由下式表示：

圖 2  P2T的總體架構(gòu)。它有4個(gè)階段，分別輸出1/4、1/8、1/16和1/32大小的特征圖。其中1/32大小的特征用于圖像分類，所有的特征用于諸多下游任務(wù)。

我們通過(guò)調(diào)整P2T基礎(chǔ)模塊的數(shù)量，就可以得到不同規(guī)模的P2T模型，如圖 3所示。

圖 3 不同設(shè)置下的P2T骨干網(wǎng)絡(luò)。

基于金子塔池化的MHSA （P-MHSA）

圖 4 P2T基本模塊架構(gòu)圖。(a) P2T模塊粗略圖；(b) 基于金字塔池化的多頭自我注意力（P-MHSA）架構(gòu)。

計(jì)算流程。 對(duì)于給定的特征，先將用不同步長(zhǎng)的池化操作下采樣到不同的大?。?/p>

可以得到不同大小的特征圖。因?yàn)檫@些特征圖是經(jīng)池化后的特征，對(duì)每個(gè)特征作相對(duì)的位置編碼（RPE）幾乎不產(chǎn)生計(jì)算量：

其中DWConv代表深度3x3卷積，即分組數(shù)等于通道數(shù)的3x3卷積。我們將這些編碼后的特征進(jìn)行拉直、特征拼接后，再分別生成計(jì)算自我注意力(self-attention)需要的query (Q)、key（K）、value（V）:

獲得 、、 后, 就可以計(jì)算自我注意力  :

其中  為  的通道數(shù), 用作尺度縮放。因?yàn)?nbsp; 和  的特征長(zhǎng)度相對(duì)  要小得多, P-MHSA 的計(jì)算量也要相對(duì) MHSA 小得多。

計(jì)算復(fù)雜度。 P-MHSA中主要由金字塔池化、自我注意力和一些線性層構(gòu)成。設(shè)N和C分別為輸入特征的長(zhǎng)度和通道數(shù)。金字塔池化所占用的計(jì)算量?jī)H為O(NC)。同時(shí)，也可以計(jì)算出P-MHSA的計(jì)算復(fù)雜度為：

其中M為經(jīng)金字塔池化采樣拼接后的特征長(zhǎng)度。對(duì)于默認(rèn)的{12,16,20,24}的池化參數(shù)，M約為，與PVT大致相同。

前饋網(wǎng)絡(luò)（Feed-forward Network, FFN）

在transformer中，前饋網(wǎng)絡(luò)（FFN）是一種基礎(chǔ)的特征增強(qiáng)模塊。然而傳統(tǒng)FFN僅用數(shù)個(gè)全連接層在特征的通道維度進(jìn)行增強(qiáng)，其并沒(méi)有考慮到圖像是2維特征。通過(guò)將2D局部特征學(xué)習(xí)引入到FFN中，可以進(jìn)一步提高場(chǎng)景理解所需要的特征表達(dá)能力。相對(duì)于普通的transformer使用的FFN，我們額外添加了一個(gè)深度3x3卷積來(lái)額外學(xué)習(xí)2D局部表征，即與MobileNetV2中的反向瓶頸模塊（IRB）一致。假設(shè)  是已轉(zhuǎn)化為 2 維的 ,，對(duì)原FFN引入2D局部特征學(xué)習(xí)：

3. 實(shí)驗(yàn)部分

我們?cè)趫D像分類、語(yǔ)義分割、目標(biāo)檢測(cè)、實(shí)例分割上與其他方法進(jìn)行了效果、計(jì)算量、運(yùn)行速度上的對(duì)比。通過(guò)下列的實(shí)驗(yàn)可以發(fā)現(xiàn)我們的方法要顯著優(yōu)于現(xiàn)有的方法，如Swin、Twins、MViT、PVTv2等。

圖像分類

圖 5 圖像分類結(jié)果。輸入大小統(tǒng)一為224x224。FPS是在RTX 2070下測(cè)試的結(jié)果。

語(yǔ)義分割

圖 6 在語(yǔ)義分割上的結(jié)果。結(jié)果基于Semantic FPN。輸入大小為512x512。所有實(shí)驗(yàn)基于batchsize 16的情況下訓(xùn)練8w次。

目標(biāo)檢測(cè)與實(shí)例分割

圖 7 目標(biāo)檢測(cè)與實(shí)例分割上的結(jié)果。圖中所有的方法都基于RetinaNet和Mask R-CNN。所有實(shí)驗(yàn)基于在coco上訓(xùn)練12 epochs后的結(jié)果（無(wú)多尺度訓(xùn)練）。

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

圖 8 金字塔池化的設(shè)置。可以發(fā)現(xiàn)，在相同下采樣系數(shù)（D Ratio）的前提下，金字塔池化要顯著優(yōu)于單尺度池化。

圖 9 相對(duì)于單尺度池化，該圖顯示的是在不同階段下應(yīng)用金字塔池化自注意力的效果?？梢园l(fā)現(xiàn)，隨著金字塔池化在更多的階段下的應(yīng)用，效果也變的越好。

圖 10 進(jìn)一步應(yīng)用相對(duì)位置編碼（RPE）、IRB和部分重疊的patch embedding（OPE）所帶來(lái)的效果。RPE可以進(jìn)一步大幅提升分類和效果，同時(shí)RPE的計(jì)算量?jī)H占網(wǎng)絡(luò)總計(jì)算量的0.1%（0.005 GFLOPS）。IRB則進(jìn)一步顯著地提升了分類和分割效果。OPE雖然進(jìn)一步提升的分類相對(duì)前兩者不多，但能顯著提升分割效果。

此外，P2T采用了更輕量的Hardswish激活，GELU與其相比，訓(xùn)練內(nèi)存占用額外多出了52%。我們也發(fā)現(xiàn)，采用Hardswish激活的P2T和采用GELU的P2T性能一致。所以，我們使用Hardswish激活，使得P2T可以在更小顯存上的顯卡上訓(xùn)練。

4. 總結(jié)

我們提出了基于金字塔池化的transformer基本單元P2T。相對(duì)于單尺度池化，金字塔池化既大幅減少了transformer中MHSA的計(jì)算量，又提供了強(qiáng)大的多層次的特征表達(dá)。我們利用P2T基本單元建立了一系列的P2T骨干：P2T-Tiny/Small/Base/Large。在圖像分類、語(yǔ)義分割、目標(biāo)檢測(cè)、實(shí)例分割等任務(wù)上，基于P2T的方法也顯著要好于其他骨干網(wǎng)絡(luò)，如ResNet, ResNeXts, Res2Nets, 以及PVT、Swin、Twins、PVTv2。

相關(guān)文獻(xiàn)

[1] An image is worth 16x16 words: Transformers for image recognition at scale, ICLR 2021

[2] Training data-efficient image transformers & distillation through attention, ICLR 2021

[3] Findings of the 2014 workshop on statistical machine translation, The Workshop on Statistical Machine Translation, 2014.

[4] ImageNet large scale visual recognition challenge, IJCV 2015

[5] Pyramid Vision Transformer: A versatile backbone for dense prediction without convolutions, ICCV 2021

[6] Swin Transformer: Hierarchical vision transformer using shifted windows, ICCV 2021

[7] Twins: Revisiting the design of spatial attention in vision transformers, NeurIPS 2021

[8] The pyramid match kernel: Discrim- inative classification with sets of image features, ICCV 2005.

[9] Beyond bags of features: Spatial pyramid matching for recognizing natural scene cate- gories, CVPR 2006.

[10] Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition, TPAMI 2015

[11] Pyramid Scene Parsing Network, CVPR 2017

本文僅做學(xué)術(shù)分享，如有侵權(quán)，請(qǐng)聯(lián)系刪文。

*博客內(nèi)容為網(wǎng)友個(gè)人發(fā)布，僅代表博主個(gè)人觀點(diǎn)，如有侵權(quán)請(qǐng)聯(lián)系工作人員刪除。