作者 | 李政

編輯 | 王曄

分享一篇最近被ICCV 2021接收的工作《Online Knowledge Distillation for Efficient Pose Estimation》，利用在線知識蒸餾進(jìn)行高效2D人體姿態(tài)估計(jì)。

論文地址： https://arxiv.org/abs/2108.02092

該篇工作由杭州師范大學(xué)和浙江大學(xué)合作完成。該文章提出了一個(gè)新的在線知識蒸餾框架OKDHP去對人體姿態(tài)估計(jì)模型進(jìn)行提升。特別地，OKDHP訓(xùn)練了一個(gè)多分支網(wǎng)絡(luò)，其中每個(gè)分支都被當(dāng)做獨(dú)立的學(xué)生模型，這里的教師不是顯式存在的，而是通過加權(quán)集成多個(gè)分支的結(jié)果后形成的集成heatmap來扮演教師的作用，通過優(yōu)化像素級別的KL Divergence損失來優(yōu)化每個(gè)學(xué)生分支模型。整個(gè)訓(xùn)練過程被簡化到了one-stage，不需要額外預(yù)訓(xùn)練的教師模型。作者在MPII和COCO上都證明了該方法的有效性。

1

引言

主流的2D姿態(tài)估計(jì)方法大多數(shù)都是基于Hourglass Network^[1](HG)，典型的工作有^[2],^[3],^[4]。其含有多個(gè)堆疊的Hourglass，通常有2-stack, 4-stack, 8-stack類型。后一個(gè)Hourglass將前一個(gè)Hourglass的結(jié)果作為輸入，不斷進(jìn)行refine，直到結(jié)尾。8-stack的結(jié)果要明顯好于4-stack，但是與之而來的問題就是計(jì)算量明顯的增加。

FPD^[5](CVPR 19')首先提出，利用傳統(tǒng)的蒸餾(KD)方法，首先訓(xùn)一個(gè)8-stack HG作為teacher，選擇一個(gè)4-stack HG作為student，然后進(jìn)行KD。參考Fig. 1。

Fig. 1 FPD (CVPR 19')

那么這篇工作明顯存在著幾點(diǎn)問題：

(1). 第一步訓(xùn)練teacher，第二步訓(xùn)student，整體是一個(gè)two-stage的過程，較為繁瑣。

(2). 如果要利用FPD訓(xùn)練一個(gè)8-stack HG的student，就需要找到一個(gè)比8-stack更高的model去作為teacher。堆疊更多的HG會存在性能收益遞減，并且?guī)碛?jì)算量直線上升。

(3). KD過程中同時(shí)使用MSE作為傳統(tǒng)任務(wù)的監(jiān)督loss和kd loss，訓(xùn)練時(shí)一個(gè)output同時(shí)針對兩個(gè)target進(jìn)行優(yōu)化會帶來明顯的沖突。

Fig.2 Network Comparison

2

方法

由以上的幾個(gè)問題就引出了我們ICCV 2021的工作，我們首先提出了利用在線知識蒸餾的方法去對網(wǎng)絡(luò)進(jìn)行訓(xùn)練。用大白話來概括一下工作的幾個(gè)核心：

(1). 我們提出了一個(gè)在線知識蒸餾的框架，即一個(gè)多分支結(jié)構(gòu)。這里的teacher不是顯式存在的，而是通過多個(gè)學(xué)生分支的結(jié)果經(jīng)過了FAU的ensemble形成的，即established on the fly，我們利用ensemble得到的結(jié)果（擁有更高的準(zhǔn)確率）來扮演teacher的角色，來KD每個(gè)的學(xué)生分支，即在Fig.2 (b)中的三個(gè)小分支。

具體來說就是，如果要得到一個(gè)4-stack HG的網(wǎng)絡(luò)，F(xiàn)PD的方式如(a)所示，先訓(xùn)練一個(gè)8-stack，然后進(jìn)行KD。而在我們的方法，如圖(b)，直接建立一個(gè)多分支網(wǎng)絡(luò)（圖中為3個(gè)分支），其中每個(gè)分支視為student，要得到一個(gè)4-stack HG，那么我們選擇在前部share 2個(gè)stack（節(jié)約計(jì)算量），后面針對每一個(gè)branch，我們將剩下的2個(gè)stack HG獨(dú)立出來，以保持diversity。三個(gè)分支產(chǎn)生的結(jié)果經(jīng)過FAU進(jìn)行ensemble，得到的ensemble heatmap再KD回每一個(gè)student分支。

我們的方法帶來的直接的好處就是，整個(gè)的KD過程簡化到One-stage，并且不需要手動的選擇一個(gè)更高performance的teacher來進(jìn)行KD。Online KD的方法直接訓(xùn)練完了之后，選擇一個(gè)最好性能的分支，去除掉其他多余分支結(jié)構(gòu)即可得到一個(gè)更高acc的目標(biāo)HG網(wǎng)絡(luò)。

那么從這里也可以直接看出我們的多分支網(wǎng)絡(luò)更省計(jì)算量，粗略的算，F(xiàn)PD的方法總共會需要8+4=12個(gè)stack參與計(jì)算，我們的方法，只會有2x4=8個(gè)stack進(jìn)行計(jì)算。

（針對diversity的問題提一下，在OKDDip^[6]中就有提及，針對這樣的一個(gè)多分支模型，每個(gè)分支之間的diversity是需要考慮的，對于每個(gè)分支，如果共享的stage過多，那么留給剩下分支的優(yōu)化空間就會被明顯縮小，分支之間在訓(xùn)練的過程中會顯式的趨于同質(zhì)化，進(jìn)而帶來的結(jié)果就是ensemble結(jié)果準(zhǔn)確率的下降。與之相反，獨(dú)立的HG數(shù)量越多也可以帶來KD性能的提升，分支數(shù)量同理，詳情請參考paper中的Table 7和8。我們在paper中將共享的HG數(shù)量設(shè)定為目標(biāo)網(wǎng)絡(luò)HG數(shù)量的一半，即目標(biāo)網(wǎng)絡(luò)8-stack，整個(gè)網(wǎng)絡(luò)就共享4-stack。）

(2). 既然是一個(gè)多分支結(jié)構(gòu)，那么每個(gè)分支的情況可不可以是adaptive的？既然在分支里更多的stack可以產(chǎn)生更好的heatmap，那么必然也就會帶來ensemble結(jié)果的提升，進(jìn)而KD的效果就會更好。于是針對(b)的這種每個(gè)分支都是一樣的balance的結(jié)構(gòu)，我們更進(jìn)一步提出了unbalance結(jié)構(gòu)。

具體的來說，要KD得到一個(gè)4-stack HG，即Fig.2 (c)中的第一個(gè)branch，2+2=4個(gè)stack的主分支，通過在輔助分支堆疊更多的HG來產(chǎn)生更好的ensemble結(jié)果，這里就是第二個(gè)分支是2+4=6個(gè)stack，第三個(gè)分支2+6=8個(gè)stack的情況。

在不考慮訓(xùn)練計(jì)算量的情況下，在部署時(shí)移除輔助分支，相比于balance結(jié)構(gòu)，可以得到更好的target student，即目標(biāo)的4-stack HG。

(3). 這里的FAU，即Feature Aggregation Unit，是用來對每個(gè)分支產(chǎn)生的結(jié)果進(jìn)行一個(gè)帶有weight的channel-wise的ensemble。即將每個(gè)heatmap按照生成的權(quán)重進(jìn)行集成。具體的結(jié)構(gòu)如Fig.3所示。

Fig.3 Feature Aggregation Unit

針對人體姿態(tài)估計(jì)這種場景下，其實(shí)是存在著很多的尺度變化問題，就比如一個(gè)人在一張圖片中，既可以是貼的很近，占滿了整張圖片，也可以是離得很遠(yuǎn)，只在圖片中占小小的一個(gè)部分。受到SKNet^[7]的啟發(fā)，在原先的3x3, 5x5的基礎(chǔ)上，我拓展出了7x7和avg pool來捕捉更大范圍的信息，進(jìn)而來生成對應(yīng)每個(gè)分支產(chǎn)生的heatmap的weight。消融實(shí)驗(yàn)證明了FAU確實(shí)是要比普通的attention方法提高更多。（小聲：avg pool這branch我做了實(shí)驗(yàn)試了一下，有一丁丁的提升，灰常?。?/span>

另外值得一提的是，我們的OKDHP方法不僅是對于hourglass類別的網(wǎng)絡(luò)有著明顯的提升，對于其他的pose estimation網(wǎng)絡(luò)也有效果，我們將其拓展到了一個(gè)非常lightweight且能夠real time進(jìn)行pose estimation的開源實(shí)現(xiàn)MobilePose^[8]上面，（在paper的末尾位置的Supplementary Material里）

這里的MobilePose可以粗略的將這個(gè)網(wǎng)絡(luò)分為encoder（backbone）和decoder部分。我們選擇去share整個(gè)的encoder部分，然后建立三個(gè)獨(dú)立的分支，每個(gè)獨(dú)立的分支都對應(yīng)一個(gè)完整的decoder部分。FAU結(jié)構(gòu)保持不變。結(jié)果如下：

可以看到，對于更加輕量化的backbone結(jié)果，我們的OKDHP方法可以獲得更明顯的漲點(diǎn)。

3

實(shí)驗(yàn)部分

在MPII validation和testing set上

在Fig. 5上，我們針對2-stack，4-stack和8-stack都進(jìn)行了實(shí)驗(yàn)，默認(rèn)的實(shí)驗(yàn)條件是3分支結(jié)構(gòu)。共享的HG數(shù)量是整個(gè)目標(biāo)網(wǎng)絡(luò)HG的一半，即要訓(xùn)練一個(gè)8-stack HG，會share 4個(gè)HG，獨(dú)立4個(gè)HG。FAU結(jié)構(gòu)不隨HG數(shù)量改變。

更進(jìn)一步對比我們的Balance和Unbalance結(jié)構(gòu)：

更為具體的Unbalance結(jié)構(gòu)的精度情況：

這里面也就對應(yīng)上了Fig. 1(c)里面的結(jié)構(gòu)，target就是一個(gè)4-stack HG的網(wǎng)絡(luò)，那么這里選擇6-stack和8-stack作為輔助分支，確實(shí)是取得了更好的結(jié)果。

在COCO val 2017上，

也有不錯(cuò)的性能提升。

Acknowledegment：

我們非常感謝浙大的陳德仿博士給出的非常nice的建議和討論。

參考

1. ^Stacked hourglass networks for human pose estimation. ECCV 16'

2. ^Multi-Context Attention for Human Pose Estimation. CVPR 17'

3. ^Learning Feature Pyramids for Human Pose Estimation. ICCV 17'

4. ^Multi-Scale Structure-Aware Network for Human Pose Estimation. ECCV 18'

5. ^Fast Human Pose Estimation. CVPR 19'

6. ^Online Knowledge Distillation with Diverse Peers. AAAI 2020

7. ^Selective kernel networks. CVPR 19'

8. ^https://github.com/YuliangXiu/MobilePose-pytorch

雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)

雷峰網(wǎng)特約稿件，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。