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