UniFormer用Transformer將3D卷積和時空自注意力有效地統一，是目前最強最快的視頻理解架構。

本文主要向大家介紹UniFormer，包括ICLR2022接受的video backbone，以及為下游密集預測任務設計的拓展版本。本套框架在各種任務上都能取得了比現有SOTA模型更好的性能。

目前UniFormer框架代碼、模型、日志以及訓練腳本都已開源，歡迎大家試用反饋！

?UniFormer代碼鏈接：

https://github.com/Sense-X/UniFormer

🪐論文標題：

UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning

🪐論文鏈接：

https://arxiv.org/abs/2201.04676

🪐擴展版本：

UniFormer: Unifying Convolution and Self-attention for Visual Recognition

https://arxiv.org/abs/2201.09450

🪐代碼鏈接：

https://github.com/Sense-X/UniFormer

“

相對SOTA，UniFormer的性能提升

1

圖像分類

在Token Labeling 的加持下，僅靠ImageNet-1K訓練。

39GFLOPs的UniFormer-L-384在ImageNet上實現了86.3%的top-1精度。

統一模型架構UniFormer

2

視頻分類

僅用ImageNet-1K預訓練，UniFormer-B在Kinetics-400和Kinetics-600上分別取得了82.9%和84.8% 的top-1精度（比使用JFT-300M預訓練，相近性能的ViViT 的GFLOPs少16倍）。

在Something-Something V1和V2上分別取得60.9%和71.2%的top-1精度，為同期模型的SOTA。

統一模型架構UniFormer

3

密集預測

僅用ImageNet-1K預訓練，COCO目標檢測任務上取得了53.8的box AP與46.4的mask AP；ADE20K語義分割任務上取得了50.8的mIoU；COCO姿態估計任務上取得了77.4的AP。后文將會介紹為下游任務設計的訓練和測試時模型適配。

點擊可看高清大圖

圖注：

圖像分類與視頻分類任務性能比較（上方為ImageNet上224x224與384x384分辨率輸入）

統一模型架構UniFormer

“

UniFormer設計靈感

對image和video上的representation learning而言，目前存在兩大痛點：

• local redundancy:
  
  視覺數據在局部空間/時間/時空鄰域具有相似性，這種局部性質容易引入大量低效的計算。

• global dependency:
  
  要實現準確的識別，需要動態地將不同區域中的目標關聯，建模長時依賴。

現有的兩大主流模型CNN和ViT，往往只關注解決以上部分問題。

Convolution只在局部小鄰域聚合上下文，天然地避免了冗余的全局計算，但受限的感受也難以建模全局依賴。

而self-attention通過比較全局相似度，自然將長距離目標關聯，但通過如下可視化我們可以發現，ViT在淺層編碼局部特征十分低效。

圖注：

DeiT可視化后，可以發現即便經過三層self-attention，輸出特征仍保留了較多的局部細節。我們任選一個token作為query，可視化attention矩陣可以發現，被關注的token集中在3x3鄰域中（紅色越深表示關注越多）

圖注：

TimeSformer可視化后，同樣可以發現即便是經過三層self-attention，輸出的每一幀特征仍保留了較多的局部細節。我們任選一個token作為query，可視化spatial attention和temporal attention矩陣都可以發現，被關注的token都只在局部鄰域中（紅色越深表示關注越多）。

無論是spatial attention抑或是temporal attention，在ViT的淺層，都僅會傾向于關注query token的鄰近token。

但是attention矩陣是通過全局token相似度計算得到的，這無疑帶來了大量不必要的計算。相較而言，convolution在提取這些淺層特征時，無論是在效果上還是計算量上都具有顯著的優勢。

那么為何不針對網絡不同層特征的差異，設計不同的特征學習算子，從而將convolution和self-attention有機地結合，物盡其用呢？

本論文中提出的UniFormer (Unified Transformer)，旨在以Transformer的風格，有機地統一convolution和self-attention，發揮二者的優勢，同時解決local redundancy和global dependency兩大問題，從而實現高效的特征學習。

“

UniFormer模型架構

圖注：

模型整體框架，標紅維度僅對video輸入作用，對image輸入都可視作1

模型整體框架如上圖所示，借鑒了CNN的層次化設計，每層包含多個Transformer風格的UniFormer block。

🪐UniFormer block：

每個UniFormer block主要由三部分組成，動態位置編碼DPE、多頭關系聚合器MHRA及Transformer必備的前饋層FFN，其中最關鍵的為多頭關系聚合器：

🪐MHRA：

與多頭注意力相似，將關系聚合器設計為多頭風格，每個頭單獨處理一組channel的信息。

每組的channel先通過線性變換生成上下文token，然后在token affinity的作用下，對上下文進行有機聚合。

基于前面的可視化觀察，可以認為在網絡的淺層，token affinity應該僅關注局部鄰域上下文，這與convolution的設計不謀而合。因此，將局部關系聚合設計為可學的參數矩陣：

🪐local MHRA：

其中為anchor token，為局部鄰域任一token，為可學參數矩陣，為二者相對位置，表明token affinity的值只與相對位置有關。

這樣local UniFormer block實際上與MobileNet block 的設計風格相似，都是PWConv-DWConv-PWConv（見原論文解析），不同的是我們引入了額外的位置編碼以及前饋層，這種特別的結合形式有效地增強了token的特征表達。

在網絡的深層，需要對整個特征空間建立長時關系，這與self-attention的思想一致，因此我們通過比較全局上下文相似度建立token affinity：

🪐global MHRA：

其中為不同的線性變換。先前的video transformer，往往采用時空分離的注意力機制 ，以減少video輸入帶來的過量點積運算，但這種分離的操作無疑割裂了token的時空關聯。

相反，我們的UniFormer在網絡的淺層采用local MHRA節省了冗余計算量，使得網絡在深層可以輕松使用聯合時空注意力，從而可以得到更具辨別性的video特征表達。

再者，與以往ViT中使用絕對位置編碼不同，我們這里采用卷積風格的動態位置編碼，使得網絡可以克服permutation-invariance的同時，對不同長度的輸入更友好。

流行的ViT往往采用絕對或者相對位置編碼 ，但絕對位置編碼在面對更大分辨率的輸入時，需要進行線性插值以及額外的參數微調，而相對位置編碼對self-attention的形式進行了修改。

為了適配不同分辨率輸入的需要，采用了最近流行的卷積位置編碼設計動態位置編碼:

🪐DPE：

其中DWConv為零填充的的深度可分離卷積。

一方面，卷積對任何輸入形式都很友好，也很容易拓展到空間維度統一編碼時空位置信息。

另一方面，深度可分離卷積十分輕量，額外的零填充可以幫助每個token確定自己的絕對位置。

“

UniFormer整體框架

🪐圖像分類：

模型細節

我們設計了三種不同規模的模型，每個模型包含4層，前兩層使用local MHRA，后兩層使用global MHRA。

對于local MHRA，卷積核大小為5x5，歸一化使用BN（使用LN性能較差）。

對于global MHRA，每個head的channel數為64，歸一化使用LN。動態位置編碼卷積核大小為3x3，FFN的拓展倍數為4。

對于特征下采樣，我們采用非重疊卷積，其中第一次下采樣卷積核大小為4x4、步長為4x4，其余三次下采樣卷積核大小為2x2、步長為2x2。

在每次下采樣卷積之后，額外增加LN歸一化。網絡最后接平均池化層與線性分類層，輸出最終預測。當使用Token Labeling時，額外加入一個線性分類層以及輔助損失函數。

對于UniFormer-S，我們設計了增強版本，每層block數量為[3, 5, 9, 3]，下采樣使用重疊卷積，FLOPs控制為4.2G，保證與其他SOTA模型可比。

🪐視頻分類：

對于video使用的3D backbone，通過加載ImageNet-1K預訓練的UniFormer-S和UniFormer-B，并進行卷積核展開。

具體來說，動態位置編碼和local MHRA分別展開為3x3x3和5x5x5卷積。

對于下采樣層，只在第一次下采樣同時壓縮時間和空間維度，而在其余三次下采樣僅壓縮空間維度，即第一次下采樣卷積核大小為3x4x4、步長為2x4x4，其余卷積核大小為1x2x2、步長為1x2x2。這樣就可以在減少計算量的同時，還能保證模型的高性能。

對于global MHRA，直接繼承相應參數，將時空token序列化進行統一處理。

🪐密集預測：

對于下游密集預測任務，直接使用原始網絡作為backbone并不合適。

因為這些任務往往輸入大分辨率圖像，比如目標檢測中輸入1333x800的圖像，使用global MHRA會帶來過多的計算量。

我們以UniFormer-S為例，統計了不同分辨率輸入時，不同操作所需的計算量。

圖注：

左邊圖片為模型整體計算量與第三/四層中MatMul運算所需計算量的結果；右邊圖片為第三層采用不同風格的MHRA所需的MatMul計算量的結果

從上圖中可以看到，第三層中MHRA所需的MatMul運算隨著分辨率的增加而急劇上升，在輸入分辨率為1008x1008時，甚至占了總運算量50%以上，而第四層僅為第三層的1/28。因此，僅對第三層的MHRA進行改進。

受先前工作的啟發，將global MHRA應用在限制的窗口內，這樣會把原本的復雜度降至，其中p為窗口大小。

然而直接應用純window化操作，不可避免地會帶來性能下降，為此我們將window和global操作結合。

每個hybrid分組中包含4個block，前3個為window block，最后1個為global block。UniFormer-S和UniFormer-B分別包含2個和5個分組。

不同任務在訓練和測試時使用特定改進

如上圖所示，在五種任務上，訓練和測試采用特定改進。

對目標檢測任務，由于訓練和測試時輸入分辨率都很大（如1333x800），在第三層都采用hybrid block。

對姿態估計任務，輸入分辨率相對較小（如384x288），在第三層采用原本的gloabl block。

而對于語義分割任務，由于往往在測試時使用幾倍于訓練輸入的分辨率（如2048x512 vs. 512x512），因此在訓練時，對第三層采用global block。

而在測試時采用hybrid block，但需要注意測試時hybrid block中window size需要與訓練時global block的感受野一致（如32x32），感受野不一致會導致急劇性能下降。

這種設計可以保證訓練高效的同時，提高測試的性能。

“

實驗結果

🪐圖像分類：

我們用ImageNet-1K進行圖像分類實驗，采用了DeiT 的代碼與訓練策略，UniFormer-S/B/L使用droppath比例為0.1/0.3/0.4，對大模型額外加入Layer Scale防止網絡訓練崩潰 。對于Token Labeling，沿用了其代碼訓練框架與超參。

ImageNet-1K

結果如上所示，其中帶*為加入Token Labaleing進行訓練，UniFormer-B\dagger為前述UniFormer-B的增強版本。

可以看到UniFormer在不同規模下都取得了SOTA性能，UniFormer-L-384僅需39G的FLOPs，即可取得86.3%的top-1精度。

🪐視頻分類：

我們在Kinetics-400/600以及Something-Something V1/V2上進行了視頻分類實驗，沿用了MViT的代碼和訓練策略。

對Kinetics采用dense采樣方式，加載了ImageNet-1K的預訓練模型加速訓練，droppath比例保持與前述一致。

對Sth-Sth采用uniform采樣方式，加載了Kinetics的預訓練，droppath設置為前述兩倍，并不進行水平翻轉。

左：Kinetics

?右：Sth-Sth

結果如上圖所示，僅使用ImageNet-1K預訓練，我們在Kinetics上取得了與使用JFT-300M預訓練的ViViT-L、使用ImageNet-21K預訓練的Swin-B相近的性能，計算量大幅減小。

而在Sth-Sth上，UniFormer大幅高于先前CNN和ViT的結果，取得了新的SOTA結果，在Sth-Sth V1上為61.2%，V2上為71.4%。

🪐目標檢測與實例分割：

我們在COCO2017上進行了目標檢測和實例分割實驗，沿用了mmdetection 的代碼框架，配置了Mask R-CNN與Cascade Mask R-CNN框架，使用Swin Transformer的訓練參數，均取得了SOTA性能。

COCO目標檢測，Mask R-CNN

COCO目標檢測，Cascade Mask R-CNN

🪐語義分割：

我們在ADE20K上進行了語義分割實驗，沿用了mmsegmentation 的代碼框架，配置了Semantic FPN與UperNet兩種框架，分別使用了PVT和Swin的訓練參數，均取得了SOTA性能。

ADE20K語義分割

左：Semantic FPN 右：UperNet

🪐人體姿態估計：

我們在COCO2017上進行了姿態估計實驗，沿用了mmpose 的代碼框架，配置了簡單的Top-down框架，使用了HRFormer的訓練參數，取得了SOTA性能。

COCO姿態估計，Top-down

🪐消融實驗：

我們進行了詳盡的消融實驗，首先在圖像和視頻分類任務上驗證了backbone的結構設計。其次，對video backbone的預訓練、訓練以及測試策略進行了探究。最后，驗證了下游密集預測任務上改進的有效性。

🤖?

Model designs for image and video backbones

結構設計

FFN：

首先將local block替換為MobileNet block，其中ReLU替換為GeLU，expand ration設置為3保持計算量一致，保留動態編碼以公平比較，可以發現local block在圖像和視頻分類任務上都會明顯優于MobileNet block。由此可知transformer風格，以及其特有的FFN確實增強了token的特征表達。

DPE：

將動態位置編碼去掉，在圖像和視頻任務上性能均下降，視頻任務上更是掉了1.7%，由此可知位置編碼有助于更好的時空特征學習。

Local MHRA size：

將local MHRA的卷積核大小設置為3、5、7、9，性能差異并不大，最終采用大小為5的卷積核，以取得最好的計算量與準確率的權衡。

MHRA configuration：

由純local MHRA（LLLL）出發，逐層替換使用global MHRA。

結果可以發現，僅使用local MHRA時，計算量很小，但性能下降明顯。逐層替換global MHRA后，性能逐漸提升。

但全部替換為global MHRA后，視頻分類準確率急劇下降，計算量急劇上升，這主要是因為網絡缺失了提取細節特征的能力，冗余的attention在有限的video數據下導致了急劇的過擬合。

🤖?

Pre-training, training and testing for video backbone

遷移性能、卷積核展開方式、訓練采樣方式

Transfer learning：

如上圖所示，表11比較了不同結構的預訓練性能以及小數據集上的遷移學習性能，可以發現，聯合的時空學習方式，不僅在預訓練模型上性能更好，在小數據集上的遷移學習性能提升明顯。

而純local MHRA以及時空分離的學習方式，遷移小數據訓練未能帶來提升。

Infalting methods：

如上圖所示，表12中比較了是否對卷積核進行展開，可以發現，展開為3D卷積核，在場景相關的數據集Kinetics-400上性能接近，但在時序相關的數據集Sth-Sth V1上提升明顯，尤其是在強的預訓練的加持下，這表明3D卷積核對時空特征的學習能力更強。

Sampling strides of dense sampling：

如上圖所示，表13中比較了我們在訓練Kinetics時，使用不同間隔采樣的結果。可以發現更稀疏的采樣，在單clip測試上效果往往更好，但在多clip測試時，間隔4幀采樣更好。

Sampling methods of Kinetics pre-trained model：

由于加載Kinetics預訓練模型訓練Sth-Sth，而Sth-Sth采用uniform采樣，所以有必要知道預訓練覆蓋更多幀是否能帶來提升。

如上圖所示，表14的結果表明，預訓練的不同采樣方式差別并不大，16x4采樣在大部分條件下性能都較好。

🤖?

左：不同數據集預訓練
右：不同測試策略? ?

Pre-trained dataset scales：

如上圖所示，圖6比較了不同規模數據預訓練的結果，可以發現對于小模型，大數據集預訓練的提升非常明顯，而對于大模型則相差無幾。

Testing strategies：

如上圖所示，圖7比較了不同的測試策略，可以發現對于使用dense采樣方式訓練的場景相關數據集Kinetics而言，多clip測試方案較好，且1x4綜合性能最優。對于使用uniform采樣方式訓練的時序相關數據集Sth-Sth而言，多crop測試方案較好，且3x1綜合性能最好。

🤖?

Adaption designs for downstream tasks

下游改進

object detection：

如上圖所示，表15比較了目標檢測任務上，第三層采用不同類型block的結果。盡管在1x訓練時，hybrid block的性能比純global block性能略差，但經過3x的充分訓練后，hybrid block的性能已經能和純global block持平。

semantic segmentation：

如上圖所示，表16比較了語義分割任務上，第三層采用不同類型block的結果。可以發現更大的窗口，以及global block的使用都能明顯提升性能，由于純global block計算量較大，我們采用性能相近的hybrid block。

pose estimation：

如上圖所示，表17分別比較了姿態估計任務上，第三層采用不同類型block的結果。由于圖像分辨率較小，zero padding消耗了更多的計算量。

🤖?

Visualizations

圖像分類與視頻分類

目標檢測、語義分割、姿態估計

“

總結與思考

在過去一年多，我們在視頻模型設計上嘗試了：

CNN（CTNet，ICLR2021）

ViT（UniFormer，ICLR2022）

MLP（MorphMLP，arxiv）

我們發現：Transformer風格的模塊+CNN的層次化架構+convolution的局部建模+DeiT強大的訓練策略，保證了模型的下限不會太低。

但相比convolution以及linear layer而言，self-attention的性價比仍是最高的，同等計算量尤其是小模型下，self-attention帶來的性能提升明顯，并且對大數據集預訓練的遷移效果更好。

?CTNet鏈接：

https://github.com/Andy1621/CT-Net

?UniFormer鏈接：

https://github.com/Sense-X/UniFormer

?MorphMLP鏈接：

https://arxiv.org/abs/2111.12527

不過傳統ViT對不同分辨率輸入并不友好，并且對大分辨率輸入的計算量難以承受，這在我們的工作UniFormer中都盡可能以簡潔的方式解決，DWConv有限制地引入，也并不會導致過多的顯存開銷與速度延遲，最后在不同的任務上都能取得很好的性能權衡。

UniFormer提供了一個盡可能簡單的框架，也希望后面的研究工作能在這個框架的基礎上，去考慮視頻中的運動信息、時空維度的冗余性、幀間的長時關系建模等等更復雜的問題，實現更大的突破。

商湯研究院基模型組火熱招聘中

實習/全職均可

技術方向：?

1. 超大模型設計與優化，大規模神經網絡理解；?

2. 基礎研究與應用，包括但不限于：目標檢測/識別/分割，知識蒸餾，基礎模型結構設計等；?

3. 無監督/半監督訓練;?

4. 大規模數據訓練優化、通用表征學習；?

5. 長尾任務，開集類別檢測等；?

6. 輕量化模型設計與優化；?

7. AutoML相關技術研發；?

團隊優勢：?

1. 1000+獨有GPU(?V100+A100)，5000+共享GPU；

2. 良好的研究氛圍與技術指導，有足夠的warm?up周期；

3. 團隊技術積累豐富，對解決具有挑戰性的問題充滿激情，獲得多項著名競賽冠軍，如ImageNet2017、OpenImage2019、ActivityNet2020、NIST?FRVT、MOT2016、MMIT2019, MFR等 。?

崗位要求：?

1. 熟練掌握機器學習（特別是深度學習）和計算機視覺的基本方法；

2. 具備以上一個或多個技術方向的研究經歷，對該領域技術理解扎實；

3. 優秀的分析問題和解決問題的能力，對解決具有挑戰性的問題充滿激情，自我驅動力強；

4. 有較強的研究能力優先，如發表過第一作者CCF A類會議或期刊等論文。?

聯系方式（簡歷投遞，請注明實習或全職）：?

songguanglu@sensetime.com

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

·

·

與50位技術專家面對面20年技術見證，附贈技術全景圖

總結

以上是生活随笔為你收集整理的比SOTA模型更全能！商汤科技和上海人工智能实验室联手打造统一模型架构UniFormer...的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。