目前,Video Pose Transformer(VPT)在基於影片的三維人體姿態估計領域取得了最領先的效能。近年來,這些 VPT 的計算量變得越來越大,這些巨大的計算量同時也限制了這個領域的進一步發展,對那些計算資源不足的研究者十分不友好。例如,訓練一個 243 幀的 VPT 模型通常需要花費好幾天的時間,嚴重拖慢了研究的進度,併成為了該領域亟待解決的一大痛點。
那麼,該如何有效地提升 VPT 的效率同時幾乎不損失精度呢?
來自北京大學的團隊提出了一種基於沙漏 Tokenizer 的高效三維人體姿態估計框架HoT,用來解決現有影片姿態 Transformer(Video Pose Transformer,VPT)高計算需求的問題。該框架可以即插即用無縫地整合到 MHFormer,MixSTE,MotionBERT 等模型中,降低模型近 40% 的計算量而不損失精度,程式碼已開源。
標題:Hourglass Tokenizer for Efficient Transformer-Based 3D Human Pose Estimation
論文地址:https://arxiv.org/abs/2311.12028
程式碼地址:https://github.com/NationalGAILab/HoT
研究動機
在 VPT 模型中,通常每一幀影片都被處理成一個獨立的 Pose Token,透過處理長達數百幀的影片序列(通常是 243 幀乃至 351 幀)來實現卓越的效能表現,並且在 Transformer 的所有層中維持全長的序列表示。然而,由於 VPT 中自注意力機制的計算複雜度與 Token 數量(即影片幀數)的平方成正比關係,當處理具有較高時序解析度的影片輸入時,這些模型不可避免地帶來了巨大的計算開銷,使得它們難以被廣泛部署到計算資源有限的實際應用中。此外,這種對整個序列的處理方式沒有有效考慮到影片序列內部幀之間的冗餘性,尤其是在視覺變化不明顯的連續幀中。這種資訊的重複不僅增加了不必要的計算負擔,而且在很大程度上並沒有對模型效能的提升做出實質性的貢獻。
因此,要想實現高效的 VPT,本文認為首先需要考慮兩個因素:
時間感受野要大:雖然直接減短輸入序列的長度能夠提升 VPT 的效率,但這樣做會縮小模型的時間感受野,進而限制模型捕獲豐富的時空資訊,對效能提升構成制約。因此,在追求高效設計策略時,維持一個較大的時間感受野對於實現精確的估計是至關重要的。
影片冗餘得去除:由於相鄰幀之間動作的相似性,影片中經常包含大量的冗餘資訊。此外,已有研究指出,在 Transformer 架構中,隨著層的加深,Token 之間的差異性越來越小。因此,可推斷出在 Transformer 的深層使用全長的 Pose Token 會引入不必要的冗餘計算,而這些冗餘計算對於最終的估計結果的貢獻有限。
基於這兩方面的觀察,作者提出對深層 Transformer 的 Pose Token 進行剪枝,以減少影片幀的冗餘性,同時提高 VPT 的整體效率。然而,這引發了一個新的挑戰:剪枝操作導致了 Token 數量的減少,這時模型不能直接估計出與原影片序列相匹配數量的三維姿態估計結果。這是因為,在傳統的 VPT 模型中,每個 Token 通常對應影片中的一幀,剪枝後剩餘的序列將不足以覆蓋原影片的全部幀,這在估計影片中所有幀的三維人體姿態時成為一個顯著的障礙。因此,為了實現高效的 VPT,還需兼顧另一個重要因素:
Seq2seq 的推理:一個實際的三維人體姿態估計系統應當能夠透過 seq2seq 的方式進行快速推理,即一次性從輸入的影片中估計出所有幀的三維人體姿態。因此,為了實現與現有 VPT 框架的無縫整合並實現快速推理,需要保證 Token 序列的完整性,即恢復出與輸入影片幀數相等的全長 Token。
基於以上三點思考,作者提出了一種基於沙漏結構的高效三維人體姿態估計框架,⏳ Hourglass Tokenizer (HoT)。總的來說,該方法有兩大亮點:
簡單的 Baseline、基於 Transformer 通用且高效的框架
HoT是第一個基於 Transformer 的高效三維人體姿態估計的即插即用框架。如下圖所示,傳統的 VPT 採用了一個 “矩形” 的正規化,即在模型的所有層中維持完整長度的 Pose Token,這帶來了高昂的計算成本及特徵冗餘。與傳統的 VPT 不同,HoT 先剪枝去除冗餘的 Token,再恢復整個序列的 Token(看起來像一個 “沙漏”),使得 Transformer 的中間層中僅保留少量的 Token,從而有效地提升了模型的效率。HoT 還展現了極高的通用性,它不僅可以無縫整合到常規的 VPT 模型中,不論是基於 seq2seq 還是 seq2frame 的 VPT,同時也能夠適配各種 Token 剪枝和恢復策略。
效率和精度兼得
HoT揭示了維持全長的姿態序列是冗餘的,使用少量代表性幀的 Pose Token 就可以同時實現高效率和高效能。與傳統的 VPT 模型相比,HoT 不僅大幅提升了處理效率,還實現了高度競爭性甚至更好的結果。例如,它可以在不犧牲效能的情況下,將 MotionBERT 的 FLOPs 降低近 50%;同時將 MixSTE 的 FLOPs 降低近 40%,而效能僅輕微下降 0.2%。
模型方法
提出的 HoT 整體框架如下圖所示。為了更有效地執行 Token 的剪枝和恢復,本文提出了 Token 剪枝聚類(Token Pruning Cluster,TPC)和 Token 恢復注意力(Token Recovering Attention,TRA)兩個模組。其中,TPC 模組動態地選擇少量具有高語義多樣性的代表性 Token,同時減輕影片幀的冗餘。TRA 模組根據所選的 Token 來恢復詳細的時空資訊,從而將網路輸出擴充套件到原始的全長時序解析度,以進行快速推理。
Token 剪枝聚類模組
本文認為選取出少量且帶有豐富資訊的 Pose Token 以進行準確的三維人體姿態估計是一個難點問題。
為了解決該問題,本文認為關鍵在於挑選那些具有高度語義多樣性的代表性 Token,因為這樣的 Token 能夠在降低影片冗餘的同時保留必要的資訊。基於這一理念,本文提出了一種簡單、有效且無需額外引數的 Token 剪枝聚類(Token Pruning Cluster,TPC)模組。該模組的核心在於鑑別並去除掉那些在語義上貢獻較小的 Token,並聚焦於那些能夠為最終的三維人體姿態估計提供關鍵資訊的 Token。透過採用聚類演算法,TPC 動態地選擇聚類中心作為代表性 Token,藉此利用聚類中心的特性來保留原始資料的豐富語義。
TPC 的結構如下圖所示,它先對輸入的 Pose Token 在空間維度上進行池化處理,隨後利用池化後 Token 的特徵相似性對輸入 Token 進行聚類,並選取聚類中心作為代表性 Token。
Token 恢復注意力模組
TPC 模組有效地減少了 Pose Token 的數量,然而,剪枝操作引起的時間解析度下降限制了 VPT 進行 seq2seq 的快速推理。因此,需要對 Token 進行恢復操作。同時,考慮到效率因素,該恢復模組應當設計得輕量級,以最小化對總體模型計算成本的影響。
為了解決上述挑戰,本文設計了一個輕量級的 Token 恢復注意力(Token Recovering Attention,TRA)模組,它能夠基於選定的 Token 恢復詳細的時空資訊。透過這種方式,由剪枝操作引起的低時間解析度得到了有效擴充套件,達到了原始完整序列的時間解析度,使得網路能夠一次性估計出所有幀的三維人體姿態序列,從而實現 seq2seq 的快速推理。
TRA 模組的結構如下圖所示,其利用最後一層 Transformer 中的代表性 Token 和初始化為零的可學習 Token,透過一個簡單的交叉注意力機制來恢復完整的 Token 序列。
應用到現有的 VPT
在討論如何將所提出的方法應用到現有的 VPT 之前,本文首先對現有的 VPT 架構進行了總結。如下圖所示,VPT 架構主要由三個組成部分構成:一個姿態嵌入模組用於編碼姿態序列的空間與時間資訊,多層 Transformer 用於學習全域性時空表徵,以及一個迴歸頭模組用於迴歸輸出三維人體姿態結果。
根據輸出的幀數不同,現有的 VPT 可分為兩種推理流程:seq2frame 和 seq2seq。在 seq2seq 流程中,輸出是輸入影片的所有幀,因此需要恢復原始的全長時序解析度。如 HoT 框架圖所示的,TPC 和 TRA 兩個模組都被嵌入到 VPT 中。在 seq2frame 流程中,輸出是影片中心幀的三維姿態。因此,在該流程下,TRA 模組是不必要的,只需在 VPT 中整合 TPC 模組即可。其框架如下圖所示。
實驗結果
消融實驗
在下表,本文給出了在 seq2seq(*)和 seq2frame(†)推理流程下的對比。結果表明,透過在現有 VPT 上應用所提出的方法,本方法能夠在保持模型引數量幾乎不變的同時,顯著減少 FLOPs,並且大幅提高了 FPS。此外,相比原始模型,所提出的方法在效能上基本持平或者能取得更好的效能。
本文還對比了不同的 Token 剪枝策略,包括注意力分數剪枝,均勻取樣,以及選擇前 k 個具有較大運動量 Token 的運動剪枝策略,可見所提出的 TPC 取得了最好的效能。
本文還對比了不同的 Token 恢復策略,包括最近鄰插值和線性插值,可見所提出的 TRA 取得了最好的效能。
與 SOTA 方法的對比
當前,在 Human3.6M 資料集上,三維人體姿態估計的領先方法均採用了基於 Transformer 的架構。為了驗證本方法的有效性,作者將其應用於三個最新的 VPT 模型:MHForme,MixSTE 和 MotionBERT,並與它們在引數量、FLOPs 和 MPJPE 上進行了比較。
如下表所示,本方法在保持原有精度的前提下,顯著降低了 SOTA VPT 模型的計算量。這些結果不僅驗證了本方法的有效性和高效率,還揭示了現有 VPT 模型中存在著計算冗餘,並且這些冗餘對最終的估計效能貢獻甚小,甚至可能導致效能下降。此外,本方法可以剔除掉這些不必要的計算量,同時達到了高度競爭力甚至更優的效能。
程式碼執行
作者還給出了 demo 執行(https://github.com/NationalGAILab/HoT),整合了 YOLOv3 人體檢測器、HRNet 二維姿態檢測器、HoT w. MixSTE 二維到三維姿態提升器。只需下載作者提供的預訓練模型,輸入一小段含有人的影片,便可一行程式碼直接輸出三維人體姿態估計的 demo。
python demo/vis.py --video sample_video.mp4
執行樣例影片得到的結果:
小結
本文針對現有 Video Pose Transforme(VPT)計算成本高的問題,提出了沙漏 Tokenizer(Hourglass Tokenizer,HoT),這是一種即插即用的 Token 剪枝和恢復框架,用於從影片中高效地進行基於 Transformer 的 3D 人體姿勢估計。研究發現,在 VPT 中維持全長姿態序列是不必要的,使用少量代表性幀的 Pose Token 即可同時實現高精度和高效率。大量實驗驗證了本方法的高度相容性和廣泛適用性。它可以輕鬆整合至各種常見的 VPT 模型中,不論是基於 seq2seq 還是 seq2frame 的 VPT,並且能夠有效地適應多種 Token 剪枝與恢復策略,展示出其巨大潛力。作者期望 HoT 能夠推動開發更強、更快的 VPT。