TPAMI 2022｜金字塔池化的骨干網絡，各大任務都漲點！南開&達摩院聯合推出P2T
加入技術交流群
作者丨吳宇寰
編輯丨極市平臺
導讀
該工作主要受PVT和MViT的啟發(fā)，第一次將金字塔池化融入到視覺transformer的骨干網絡中，在減小輸入序列長度的同時去捕捉多層次的豐富特征表達。P2T在圖像分類、語義分割、目標檢測、實例分割等多個領域中，并取得了比現有CNN/Transformer骨干更優(yōu)異的性能。
作者列表：吳宇寰，劉云，占新，程明明
作者單位：南開大學，阿里達摩院
代碼位置：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技術在計算機視覺領域的興起近年來，NLP中的transformer技術在計算機視覺領域掀起新的熱潮。Transformer可以充分地捕捉長距離的全局特征關系，正好滿足了計算機視覺大感受野的需求。視覺transformer [1] 也在近年來的各大領域取得了巨大的成功，證明了純transformer架構在巨大的數據規(guī)模下訓練也能在大規(guī)模的ImageNet分類數據集上與CNN方法取得相同的分類效果。DeiT [2] 則隨后證明了使用知識蒸餾可以只在ImgeNet-1K分類數據集上就能與CNN取得類似的效果。
然而與NLP領域特征長度有限不同，計算機視覺的圖像具有密集特性，其拉直(flattened)后的特征長度要遠遠大于NLP任務中的序列長度。比如如WMT2014數據集 [3] 的平均序列長度為25，ImageNet-1K數據集 [4] 的一般使用圖像分辨率為224x224的大小，即使使用常規(guī)的倍降采樣，拉直后的圖像序列長度仍有3136。而Transformer因其需要捕捉充分的全局注意力關系，其多頭自我注意力（multi-head self-attention, MHSA）的計算復雜度是與輸入特征呈平方關系。CNN網絡因為主要捕捉局部信息，其計算復雜度只與圖像的特征長度呈線性相關。
因此，在計算機視覺領域中應用transformer技術的關鍵困難在于如何使計算注意力關系的過程更加高效。視覺transformer [1]作為第一個解決方案，首先提出使用大的降采樣尺度（如16、32倍）來對輸入特征進行降采樣，并設計了一個純transformer的網絡骨干（backbone），并在圖像分類任務上取得了巨大的成功。而隨后PVT [5]、Swin [6]、MViT [7] 等金字塔結構的網絡骨干使用類似CNN的逐漸下采樣的手段，來支持各類下游視覺任務。
目前transformer網絡骨干的思路如上文所述，transformer應用到計算機視覺領域的關鍵問題在于如何使計算MHSA的過程更加高效。PVT和MViT對特征圖采用了單個池化操作來減少MHSA的計算量和顯存占用，相應地是進行輸入特征圖元素與不同區(qū)域之間的關系建模，而不是特征圖元素之間的關系建模。Swin transformer則是另一種思路：它先將特征劃分為小的區(qū)塊，并在每個小區(qū)塊內直接計算MHSA，即進行小區(qū)塊內元素之間的關系建模；隨后，它將區(qū)塊劃分進行偏移，同時計算偏移后的小區(qū)塊的MHSA。這樣，類似CNN網絡一樣，隨著網絡的堆疊，Swin transformer網絡的感受野會逐步增強。然而，視覺transformer的最根本特性源于其對特征全局的關系建模，這也是為什么我們要嘗試將transformer引入到過度由CNN統(tǒng)治的計算機視覺領域中。
本文的工作我們的工作主要受PVT和MViT的啟發(fā)，即使用特征降采樣的形式來減少MHSA的計算量。PVT和MViT使用單個尺度的降采樣后的特征來計算MHSA，這種特征表達的模式有限。我們想到，如果使用金字塔池化（pyramid pooling）來生成特征，既能夠大大縮小特征序列長度，又能引入更豐富的多層次特征。
金字塔池化在計算機視覺中有著悠久的歷史 [8-9]。它利用多個不同步長的池化操作來提取豐富不同感受野的池化特征，其在目標檢測 [10] 和語義分割中 [11] 也被證明十分有效。然而，它們都依賴于一個CNN骨干網絡，所以它們也只是被專門設計用于某個特定的任務。目前并沒有探索金字塔池化在骨干網絡設計的工作。而一個新型的骨干網絡可以應用于多種下游任務中。我們第一次將金字塔池化融入到視覺transformer的骨干網絡中，在減小輸入序列長度的同時去捕捉多層次的豐富特征表達。
利用金字塔池化，我們設計了基于的金字塔池化的MHSA（pooling-based MHSA, P-MHSA）。金字塔池化作為計算P-MHSA中降低序列長度的基本操作。使用P-MHSA替換掉MHSA后，即為P2T基礎模塊。通過堆疊不同數量的P2T基礎模塊，我們設計了P2T-Tiny/Small/Base/Large骨干網絡，分別對應ResNet-18/50/101/152。P2T系列的多層次特征捕捉能力也要強于之前的PVT與MViT方法。我們也在圖像分類、語義分割、目標檢測、實例分割等多個領域中，并取得了比現有CNN/Transformer骨干更優(yōu)異的性能，如ResNet, ResNeXt, Res2Net, PVT, Swin, Twins, 及PVTv2。
圖 1 P2T與其他方法在ADE20K數據集上的對比圖（參數量、mIoU精度）。
2. P2T 方法P2T的總體結構如圖 2所示。每個階段開始時有一個patch embedding模塊，用于對輸入的特征進行下采樣。第一個Patch Embedding模塊通過步長為4的7x7 卷積將輸入特征下采樣為原來的1/4，其他Patch Embedding則是將輸入特征下采樣為原來的2倍。每個Patch Embedding模塊之后則堆疊了不同數量的P2T基礎模塊。其中，P2T基本模塊的計算可由下式表示：
圖 2 P2T的總體架構。它有4個階段，分別輸出1/4、1/8、1/16和1/32大小的特征圖。其中1/32大小的特征用于圖像分類，所有的特征用于諸多下游任務。
我們通過調整P2T基礎模塊的數量，就可以得到不同規(guī)模的P2T模型，如圖 3所示。
圖 3 不同設置下的P2T骨干網絡。
基于金子塔池化的MHSA （P-MHSA）圖 4 P2T基本模塊架構圖。(a) P2T模塊粗略圖；(b) 基于金字塔池化的多頭自我注意力（P-MHSA）架構。
計算流程。 對于給定的特征，先將用不同步長的池化操作下采樣到不同的大小：
可以得到不同大小的特征圖。因為這些特征圖是經池化后的特征，對每個特征作相對的位置編碼（RPE）幾乎不產生計算量：
其中DWConv代表深度3x3卷積，即分組數等于通道數的3x3卷積。我們將這些編碼后的特征進行拉直、特征拼接后，再分別生成計算自我注意力(self-attention)需要的query (Q)、key（K）、value（V）:
獲得 、、 后, 就可以計算自我注意力 :
其中 為 的通道數, 用作尺度縮放。因為 和 的特征長度相對 要小得多, P-MHSA 的計算量也要相對 MHSA 小得多。
計算復雜度。 P-MHSA中主要由金字塔池化、自我注意力和一些線性層構成。設N和C分別為輸入特征的長度和通道數。金字塔池化所占用的計算量僅為O(NC)。同時，也可以計算出P-MHSA的計算復雜度為：
其中M為經金字塔池化采樣拼接后的特征長度。對于默認的{12,16,20,24}的池化參數，M約為，與PVT大致相同。
前饋網絡（Feed-forward Network, FFN）在transformer中，前饋網絡（FFN）是一種基礎的特征增強模塊。然而傳統(tǒng)FFN僅用數個全連接層在特征的通道維度進行增強，其并沒有考慮到圖像是2維特征。通過將2D局部特征學習引入到FFN中，可以進一步提高場景理解所需要的特征表達能力。相對于普通的transformer使用的FFN，我們額外添加了一個深度3x3卷積來額外學習2D局部表征，即與MobileNetV2中的反向瓶頸模塊（IRB）一致。假設 是已轉化為 2 維的 ,，對原FFN引入2D局部特征學習：
我們在圖像分類、語義分割、目標檢測、實例分割上與其他方法進行了效果、計算量、運行速度上的對比。通過下列的實驗可以發(fā)現我們的方法要顯著優(yōu)于現有的方法，如Swin、Twins、MViT、PVTv2等。
圖像分類圖 5 圖像分類結果。輸入大小統(tǒng)一為224x224。FPS是在RTX 2070下測試的結果。
語義分割圖 6 在語義分割上的結果。結果基于Semantic FPN。輸入大小為512x512。所有實驗基于batchsize 16的情況下訓練8w次。
目標檢測與實例分割圖 7 目標檢測與實例分割上的結果。圖中所有的方法都基于RetinaNet和Mask R-CNN。所有實驗基于在coco上訓練12 epochs后的結果（無多尺度訓練）。
消融實驗圖 8 金字塔池化的設置。可以發(fā)現，在相同下采樣系數（D Ratio）的前提下，金字塔池化要顯著優(yōu)于單尺度池化。
圖 9 相對于單尺度池化，該圖顯示的是在不同階段下應用金字塔池化自注意力的效果?？梢园l(fā)現，隨著金字塔池化在更多的階段下的應用，效果也變的越好。
圖 10 進一步應用相對位置編碼（RPE）、IRB和部分重疊的patch embedding（OPE）所帶來的效果。RPE可以進一步大幅提升分類和效果，同時RPE的計算量僅占網絡總計算量的0.1%（0.005 GFLOPS）。IRB則進一步顯著地提升了分類和分割效果。OPE雖然進一步提升的分類相對前兩者不多，但能顯著提升分割效果。
此外，P2T采用了更輕量的Hardswish激活，GELU與其相比，訓練內存占用額外多出了52%。我們也發(fā)現，采用Hardswish激活的P2T和采用GELU的P2T性能一致。所以，我們使用Hardswish激活，使得P2T可以在更小顯存上的顯卡上訓練。
4. 總結我們提出了基于金字塔池化的transformer基本單元P2T。相對于單尺度池化，金字塔池化既大幅減少了transformer中MHSA的計算量，又提供了強大的多層次的特征表達。我們利用P2T基本單元建立了一系列的P2T骨干：P2T-Tiny/Small/Base/Large。在圖像分類、語義分割、目標檢測、實例分割等任務上，基于P2T的方法也顯著要好于其他骨干網絡，如ResNet, ResNeXts, Res2Nets, 以及PVT、Swin、Twins、PVTv2。
相關文獻
[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
本文僅做學術分享，如有侵權，請聯系刪文。
*博客內容為網友個人發(fā)布，僅代表博主個人觀點，如有侵權請聯系工作人員刪除。
