作者丨happy
編輯丨極市平台
極市導讀
CNN與Transformer互相借鑒是必然趨勢，但如何借鑒并揚長避短仍需進一步深入挖掘。本文的UniFormer提供了一個非常不錯的思路，它将卷積與自注意力以transformer方式進行了統一建構UniFormer子產品，并由此建構了UniFormer。最後，作者在不同視覺任務(包含圖像分類、視訊分類、目标檢測、執行個體分割、語義分割、姿态估計)上驗證了UniFormer的超優異特性，真可謂“一力降十會”。>>加入極市CV技術交流群，走在計算機視覺的最前沿
論文連結：https://arxiv.org/abs/2201.09450
代碼連結：https://github.com/Sense-X/UniFormer
Abstract
圖像/視訊資料中的局部備援與複雜全局依賴關系使得從中學習具有判别能力的特征表達極具挑戰性。CNN與ViTs(Vision Transformers)是兩種主流的架構，CNN通過卷積有效的降低了局部備援但有限的感受野使其無法捕獲全局依賴關系，而ViT憑借自注意力可以捕獲長距離依賴，但盲相似性比對會導緻過高的備援。
為解決上述問題，我們提出一種新的UniFormer(Unified transFormer)，它能夠将卷積與自注意力的優點通過transformer進行無縫內建。不同于經典的Transformer子產品，UniFormer子產品的相關性聚合在淺層與深層分别武裝了局部全局token，解決了高效表達學習的備援與依賴問題。
基于UniFormer子產品，我們建構了一種新強力骨幹并用于不同的視覺任務，包含圖像與視訊，分類與稠密預測。無需額外你訓練資料，UniFormer在ImageNet資料及上取得了86.3%的精度；僅需ImageNet-1K預訓練，它在諸多下遊任務上取得了SOTA性能，比如Kinetics-400/600資料集上的82.9%/84.8%、Something-Something V1/V2資料集上的60.9%/71.2%、COCO檢測任務上的53.8boxAP與46.4MaskAP、ADE20K分割任務上的50.8mIoU、COCO姿态估計上的77.4AP。
Method
上圖給出了本文所提UniFormer架構示意圖，為簡單起見，我們以T幀的視訊輸入作為示例。注：圖中标紅的次元僅作用于視訊輸入，其他次元對于圖像和視訊輸入相同。
具體來說，UniFormer子產品包含三個關鍵子產品：
- Dynamic Position Embedding，DPE
- Multi-Head Relation Aggregator，MHRA
- Feed-Forward Network，FFN
對于輸入，我們首先引入DPE将位置資訊動态內建進所有tokens，它适用于任意輸入分辨率，可充分利用token順序進行更好視覺識别，可描述如下：
然後，我們采用MHRA通過探索上下文token與相關性學習對每個token進行增強，描述如下：
最後，類似ViTs，我們添加FFN進一步增強token，可描述如下：
Multi-head Relation Attention
正如前面所提到：CNN與ViTs分别聚焦于解決局部備援與全局依賴，導緻了次優性能、不必要的計算備援。為解決上述問題，我們引入了一種廣義相關性聚合(Relation Aggregator, RA)，它将卷積與自注意力統一為token相關學習。通過在淺層與深層設計局部與全局token affinity，它能夠取得更高效&有效的表達學習能力。具體來說，MHRA以多頭方式探索token相關性：
對于輸入，我們首先将其reshape為token序列。表示RA的第n個頭，表示可學習參數矩陣用于N個頭聚合。每個RA包含token上下文編碼與token親和學習。我們通過線性變換将原始token編碼為上下文token，然後RA可以通過token相關性矩陣An對token進行上下文資訊聚合。
Local MHRA
如上圖所示，盡管已有ViTs在所有token之間比較相似性，但他們最終學習了局部表達。這種自注意力備援設計帶來巨大的計算消耗。除了該發現外，我們建議在近鄰之間學習token相關性，這與卷積濾波器的設計相似。是以，我們在淺層設計了局部相關性參數矩陣。具體來說，給定輸入token ，局部RA在小範圍内進行token間相關性計算：
由于淺層的視覺内容近鄰變化很小，是以沒有必要讓上述相關性矩陣存在資料依賴性。是以，我們采用可學習參數矩陣描述上述局部token相關性，它僅依賴于相對位置資訊。
Global MHRA
在深層，長距離相關性探索非常重要，它具有與自注意力相似的思想。是以，我們從全局視角設計了token相關性矩陣：
Dynamic Position Embedding
位置資訊對于描述視覺表達非常重要。已有ViTs通過絕對/相對位置嵌入方式進行編碼，但均存在一定的不靈活性。為改善靈活性，我們采用了如下動态位置嵌入：
其中DWConv表示zero-padding深度卷積。該設計主要基于以下三點考量：
- 深度卷積對于任意輸入分辨率友好；
- 深度卷積極為輕量，是計算量-均衡均衡的重要因子；
- zero-padding有助于token具有絕對位置感覺性。
Framework
接下來，我們将針對不同下遊任務進行架構設計，包含圖像分類、視訊分類、目标檢測、語義分割、人體姿态估計等。相關架構示意圖可參見上圖。
Image Classification
前面的Figure3給出了用于圖像/視訊分類的架構示意圖，它包含四個階段。具體來說，我們在前兩個階段采用局部UniFormer子產品以降低計算備援；在後兩個階段采用全局UniFormer子產品以學習長距離token依賴。
對于局部UniFormer子產品，MHRA配置為PWConv-DWConv-PWConv(其中DWConv的尺寸為)；對于全局UniFormer子產品，MHRA配置為多頭自注意力。對于兩種UniFormer，DPE均為DWConv，尺寸為；FFN的擴充比例為4。
此外，我們對卷積使用BN，對自注意力使用LN。對于特征下采樣，我們在第一階段使用尺寸和stride均為的卷積；對于其他卷積則采用尺寸和stride為的卷積。除此之外，下采樣卷積後接LN。最後，采用GAP與全連接配接層進行分類。為滿足不同計算需求，我們設計了三種複雜度的模型，見下表。
Video Classification
基于前述圖像分類的2D骨幹，我們可以輕易的将其擴充為用于視訊分類的3D骨幹。不失一般性，我們調整Small與Base進行空時模組化。具體來說，模型架構仍保持四個階段不變。不同之處在于：所有2D卷積替換為3D卷積。DPE與局部MHRA中的DWConv濾波器尺寸為。特别的，我們在第一階段之前需要進行空時次元下采樣，此時的卷積濾波器與stride分别為。對于其他階段，我們僅僅進行空間次元下采樣，是以其他階段的下采樣卷積濾波器次元為。
在全局UniFormer子產品中，我們從3D視角采用空時注意力學習token相關性。而已有Video Transformer則進行空域與時序拆分以降低計算量、緩解過拟合，這無疑會弱化空時相關性。此外，由于所提局部UniFormer子產品可以極大節省計算量，故所提所提UniFormer可以更高效且有效的進行視訊表達學習。
Dense Prediction
稠密預測任務有助于驗證所提識别骨幹架構的泛化性。是以，我們将UniFormer骨幹使用多個主流稠密任務，包含目标檢測、執行個體分割、語義分割以及人體姿态估計。
由于大多稠密預測的輸入分辨率比較高(比如COCO檢測上的)，直接使用該骨幹會導緻具體計算量。為此，我們對不同下遊任務調整全局UniFormer子產品。上圖給出了分辨率與計算量之間的關系，很明顯：階段3的RA占據了大量的計算量，甚至高達50%，而階段4的結算量僅為階段3的1/28。是以，我們主要聚焦于修改階段3的RA。
受啟發于Swin Transformer，我們在預定于視窗(而非全局圖像範圍)内執行全局MHRA。這種處理方式可以大幅降低計算量，但它不可避免會降低模型性能。為彌補該差距，我們在階段即同時內建視窗形式與全局形式UniFormer，每個混合組包含三個視窗子產品與1個全局子產品。
基于上述設計，我們将引入用于不同稠密任務的定制骨幹：
- 目标檢測與執行個體分割：在階段3采用混合子產品；
- 姿态估計：由于輸入較小，仍采用全局子產品；
- 語義分割：由于測試時分辨率更大，故訓練時采用全局子產品，測試時采用混合子產品。這種設計可以保持訓練效率，同時可以提升測試性能。
Experiments
關于訓練細節方面資訊，我們就直接略過，直接上結果咯。
Image Classification
上表給出了ImageNet資料上的性能對比，從中可以看到：
- 在不同計算限制下，所提UniFormer均超越了其他CNN與ViTs；
- UniFormer-S取得了83.4%的精度且僅需4.2G FLOPs，分别以3.4%、2.1%、0.7%、1.8%超越了RegNetY-4G、Swin-T、CSwin-T以及CoAtNet；
- 引入了Token Labeling機制後，所提方案性能進一步提升到了86.3%，與VOLO性能相同且計算量少43%。
Video Classification
上表給出了Kinetics-400&600資料集上的性能，可以看到：
- 相比SlowFast，Uniformer-S均取得了1%名額提升且計算量少42x；
- 相比MoViNet，所提方案能取得了0.5%性能提升(82.0% vs 81.5%)，同時輸入幀數更少(16fx4 vs 120f)；
- 僅用ImageNet-1K預訓練，UniFormer-B超越了大部分現有采用更大資料集預訓練的性能。比如相比ViViT-L(JFT-300M預訓練)、Swin-B(ImageNet-21K預訓練)，UniFormerB取得了相當的性能，而計算量在兩個資料集上分别少16.7x和3.3x。
上表為Something-Something資料集上的性能對比，可以看到：
- UniFormer-S僅需42GFLOPs取得了54.4%/65.0%的優秀名額；
- 最佳模型UniFormer-B取得了61.0%/71.2%的SOTA名額。
Object Detection & Instance Segmentation
上表給出了COCO檢測與分割任務上的性能對比，基礎架構為Mask R-CNN，可以看到：
- UniFormer取得了優于所有CNN與ViTs的性能；
- 相比ResNet，所提UniFormer取得了7.0-7.6box mAP與6.7-7.2mask mAP名額提升；
- 相比SwinT，所提UniFormer取得了2.6-3.4box mAP與2.2-2.5mask mAP名額提升；
- 當采用更好訓練機制時，UniFormer-B以0.3box mAP和0.3mask mAP超越了CSwin-S、Swin-B、Focal-B等方案。
Semantic Segmentation
上表給出了ADE20K資料集上的性能對比，可以看到：
- 基于SemanticFPN架構時，相比SwinT，UniFormer-S/B取得了4.7/2.5mIoU名額提升；
- 基于UperNet架構時，UniFormer的性能提升2.5/1.9mIoU、2.7/1.2MS mIoU。
Pose Estimation
上表給出了COCO姿态估計任務上的性能對比，可以看到：
- 相比SOTA CNN方案，所提UniFormer與0.4%AP名額差偶爾了HRNet-W48，同時參數量與FLOPs更低；
- 相比目前最佳HRFormer，UniFormer-B以0.2%AP名額超出，同時FLOPs更低。
沒有圖示的paper是沒有靈魂的，最後就補充個圖示效果以供參考。
—版權聲明—
—THE END—