0
本文作者: 我在思考中 | 2021-08-20 14:45 |
分享一篇最近被ICCV 2021接收的工作《Online Knowledge Distillation for Efficient Pose Estimation》,利用在線知識蒸餾進行高效2D人體姿態(tài)估計。
論文地址: https://arxiv.org/abs/2108.02092
該篇工作由杭州師范大學和浙江大學合作完成。該文章提出了一個新的在線知識蒸餾框架OKDHP去對人體姿態(tài)估計模型進行提升。特別地,OKDHP訓練了一個多分支網(wǎng)絡,其中每個分支都被當做獨立的學生模型,這里的教師不是顯式存在的,而是通過加權集成多個分支的結果后形成的集成heatmap來扮演教師的作用,通過優(yōu)化像素級別的KL Divergence損失來優(yōu)化每個學生分支模型。整個訓練過程被簡化到了one-stage,不需要額外預訓練的教師模型。作者在MPII和COCO上都證明了該方法的有效性。
1
FPD[5](CVPR 19')首先提出,利用傳統(tǒng)的蒸餾(KD)方法,首先訓一個8-stack HG作為teacher,選擇一個4-stack HG作為student,然后進行KD。參考Fig. 1。
Fig. 1 FPD (CVPR 19')
那么這篇工作明顯存在著幾點問題:
(1). 第一步訓練teacher,第二步訓student,整體是一個two-stage的過程,較為繁瑣。
(2). 如果要利用FPD訓練一個8-stack HG的student,就需要找到一個比8-stack更高的model去作為teacher。堆疊更多的HG會存在性能收益遞減,并且?guī)碛嬎懔恐本€上升。
(3). KD過程中同時使用MSE作為傳統(tǒng)任務的監(jiān)督loss和kd loss,訓練時一個output同時針對兩個target進行優(yōu)化會帶來明顯的沖突。
Fig.2 Network Comparison
2
(1). 我們提出了一個在線知識蒸餾的框架,即一個多分支結構。這里的teacher不是顯式存在的,而是通過多個學生分支的結果經(jīng)過了FAU的ensemble形成的,即established on the fly,我們利用ensemble得到的結果(擁有更高的準確率)來扮演teacher的角色,來KD每個的學生分支,即在Fig.2 (b)中的三個小分支。
具體來說就是,如果要得到一個4-stack HG的網(wǎng)絡,F(xiàn)PD的方式如(a)所示,先訓練一個8-stack,然后進行KD。而在我們的方法,如圖(b),直接建立一個多分支網(wǎng)絡(圖中為3個分支),其中每個分支視為student,要得到一個4-stack HG,那么我們選擇在前部share 2個stack(節(jié)約計算量),后面針對每一個branch,我們將剩下的2個stack HG獨立出來,以保持diversity。三個分支產(chǎn)生的結果經(jīng)過FAU進行ensemble,得到的ensemble heatmap再KD回每一個student分支。
我們的方法帶來的直接的好處就是,整個的KD過程簡化到One-stage,并且不需要手動的選擇一個更高performance的teacher來進行KD。Online KD的方法直接訓練完了之后,選擇一個最好性能的分支,去除掉其他多余分支結構即可得到一個更高acc的目標HG網(wǎng)絡。
那么從這里也可以直接看出我們的多分支網(wǎng)絡更省計算量,粗略的算,F(xiàn)PD的方法總共會需要8+4=12個stack參與計算,我們的方法,只會有2x4=8個stack進行計算。
(針對diversity的問題提一下,在OKDDip[6]中就有提及,針對這樣的一個多分支模型,每個分支之間的diversity是需要考慮的,對于每個分支,如果共享的stage過多,那么留給剩下分支的優(yōu)化空間就會被明顯縮小,分支之間在訓練的過程中會顯式的趨于同質(zhì)化,進而帶來的結果就是ensemble結果準確率的下降。與之相反,獨立的HG數(shù)量越多也可以帶來KD性能的提升,分支數(shù)量同理,詳情請參考paper中的Table 7和8。我們在paper中將共享的HG數(shù)量設定為目標網(wǎng)絡HG數(shù)量的一半,即目標網(wǎng)絡8-stack,整個網(wǎng)絡就共享4-stack。)
(2). 既然是一個多分支結構,那么每個分支的情況可不可以是adaptive的?既然在分支里更多的stack可以產(chǎn)生更好的heatmap,那么必然也就會帶來ensemble結果的提升,進而KD的效果就會更好。于是針對(b)的這種每個分支都是一樣的balance的結構,我們更進一步提出了unbalance結構。
具體的來說,要KD得到一個4-stack HG,即Fig.2 (c)中的第一個branch,2+2=4個stack的主分支,通過在輔助分支堆疊更多的HG來產(chǎn)生更好的ensemble結果,這里就是第二個分支是2+4=6個stack,第三個分支2+6=8個stack的情況。
在不考慮訓練計算量的情況下,在部署時移除輔助分支,相比于balance結構,可以得到更好的target student,即目標的4-stack HG。
(3). 這里的FAU,即Feature Aggregation Unit,是用來對每個分支產(chǎn)生的結果進行一個帶有weight的channel-wise的ensemble。即將每個heatmap按照生成的權重進行集成。具體的結構如Fig.3所示。
Fig.3 Feature Aggregation Unit
針對人體姿態(tài)估計這種場景下,其實是存在著很多的尺度變化問題,就比如一個人在一張圖片中,既可以是貼的很近,占滿了整張圖片,也可以是離得很遠,只在圖片中占小小的一個部分。受到SKNet[7]的啟發(fā),在原先的3x3, 5x5的基礎上,我拓展出了7x7和avg pool來捕捉更大范圍的信息,進而來生成對應每個分支產(chǎn)生的heatmap的weight。消融實驗證明了FAU確實是要比普通的attention方法提高更多。(小聲:avg pool這branch我做了實驗試了一下,有一丁丁的提升,灰常小)
另外值得一提的是,我們的OKDHP方法不僅是對于hourglass類別的網(wǎng)絡有著明顯的提升,對于其他的pose estimation網(wǎng)絡也有效果,我們將其拓展到了一個非常lightweight且能夠real time進行pose estimation的開源實現(xiàn)MobilePose[8]上面,(在paper的末尾位置的Supplementary Material里)
這里的MobilePose可以粗略的將這個網(wǎng)絡分為encoder(backbone)和decoder部分。我們選擇去share整個的encoder部分,然后建立三個獨立的分支,每個獨立的分支都對應一個完整的decoder部分。FAU結構保持不變。結果如下:
可以看到,對于更加輕量化的backbone結果,我們的OKDHP方法可以獲得更明顯的漲點。
3
在Fig. 5上,我們針對2-stack,4-stack和8-stack都進行了實驗,默認的實驗條件是3分支結構。共享的HG數(shù)量是整個目標網(wǎng)絡HG的一半,即要訓練一個8-stack HG,會share 4個HG,獨立4個HG。FAU結構不隨HG數(shù)量改變。
更進一步對比我們的Balance和Unbalance結構:
更為具體的Unbalance結構的精度情況:
這里面也就對應上了Fig. 1(c)里面的結構,target就是一個4-stack HG的網(wǎng)絡,那么這里選擇6-stack和8-stack作為輔助分支,確實是取得了更好的結果。
在COCO val 2017上,
也有不錯的性能提升。
Acknowledegment:
我們非常感謝浙大的陳德仿博士給出的非常nice的建議和討論。
^Stacked hourglass networks for human pose estimation. ECCV 16'
^Multi-Context Attention for Human Pose Estimation. CVPR 17'
^Learning Feature Pyramids for Human Pose Estimation. ICCV 17'
^Multi-Scale Structure-Aware Network for Human Pose Estimation. ECCV 18'
^Fast Human Pose Estimation. CVPR 19'
^Online Knowledge Distillation with Diverse Peers. AAAI 2020
^Selective kernel networks. CVPR 19'
^https://github.com/YuliangXiu/MobilePose-pytorch
雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)
雷峰網(wǎng)特約稿件,未經(jīng)授權禁止轉載。詳情見轉載須知。