ECCV 2022 | CMU提出FKD：用于視覺識別的快速知識蒸餾框架！訓(xùn)練加速30%！

今天介紹一篇來自卡耐基梅隆大學(xué)等單位 ECCV 2022 的一篇關(guān)于快速知識蒸餾的文章，用基本的訓(xùn)練參數(shù)配置就可以把 ResNet-50 在 ImageNet-1K 從頭開始 (from scratch) 訓(xùn)練到 80.1% (不使用 mixup，cutmix 等數(shù)據(jù)增強(qiáng))，訓(xùn)練速度（尤其是數(shù)據(jù)讀取開銷）相比傳統(tǒng)分類框架節(jié)省 16% 以上，比之前 SOTA 算法快 30% 以上，是目前精度和速度雙雙最優(yōu)的知識蒸餾策略之一，代碼和模型已全部開源！

A Fast Knowledge Distillation Framework for Visual Recognition

論文和項目網(wǎng)址：http://zhiqiangshen.com/projects/FKD/index.html

代碼：https://github.com/szq0214/FKD

知識蒸餾（KD）自從 2015 年由 Geoffrey Hinton 等人提出之后，在模型壓縮，視覺分類檢測等領(lǐng)域產(chǎn)生了巨大影響，后續(xù)產(chǎn)生了無數(shù)相關(guān)變種和擴(kuò)展版本，但是大體上可以分為以下幾類：vanilla KD，online KD，teacher-free KD 等。最近不少研究表明，一個最簡單、樸素的知識蒸餾策略就可以獲得巨大的性能提升，精度甚至高于很多復(fù)雜的 KD 算法。但是 vanilla KD 有一個不可避免的缺點：每次 iteration 都需要把訓(xùn)練樣本輸入 teacher 前向傳播產(chǎn)生軟標(biāo)簽 (soft label)，這樣就導(dǎo)致很大一部分計算開銷花費在了遍歷 teacher 模型上面，然而 teacher 的規(guī)模通常會比 student 大很多，同時 teacher 的權(quán)重在訓(xùn)練過程中都是固定的，這樣就導(dǎo)致整個知識蒸餾框架學(xué)習(xí)效率很低。 針對這個問題，本文首先分析了為何沒法直接為每張輸入圖片產(chǎn)生單個軟標(biāo)簽向量然后在不同 iterations 訓(xùn)練過程中復(fù)用這個標(biāo)簽，其根本原因在于視覺領(lǐng)域模型訓(xùn)練過程數(shù)據(jù)增強(qiáng)的使用，尤其是 random-resize-cropping 這個圖像增強(qiáng)策略，導(dǎo)致不同 iteration 產(chǎn)生的輸入樣本即使來源于同一張圖片也可能來自不同區(qū)域的采樣，導(dǎo)致該樣本跟單個軟標(biāo)簽向量在不同 iterations 沒法很好的匹配。本文基于此，提出了一個快速知識蒸餾的設(shè)計，通過特定的編碼方式來處理需要的參數(shù)，繼而進(jìn)一步存儲復(fù)用軟標(biāo)簽（soft label），與此同時，使用分配區(qū)域坐標(biāo)的策略來訓(xùn)練目標(biāo)網(wǎng)絡(luò)。通過這種策略，整個訓(xùn)練過程可以做到顯式的 teacher-free，該方法的特點是既快（16%/30% 以上訓(xùn)練加速，對于集群上數(shù)據(jù)讀取緩慢的缺點尤其友好），又好（使用 ResNet-50 在 ImageNet-1K 上不使用額外數(shù)據(jù)增強(qiáng)可以達(dá)到 80.1% 的精度）。 首先我們來回顧一下普通的知識蒸餾結(jié)構(gòu)是如何工作的，如下圖所示：

知識蒸餾框架包含了一個預(yù)訓(xùn)練好的 teacher 模型（蒸餾過程權(quán)重固定），和一個待學(xué)習(xí)的 student 模型, teacher 用來產(chǎn)生 soft 的 label 用于監(jiān)督 student 的學(xué)習(xí)。可以看到，這個框架存在一個比較明顯的缺點：當(dāng) teacher 結(jié)構(gòu)大于 student 的時候，訓(xùn)練圖像前饋產(chǎn)生的計算開銷已經(jīng)超過 student，然而 teacher 權(quán)重并不是我們學(xué)習(xí)的目標(biāo)，導(dǎo)致這種計算開銷本質(zhì)上是 “無用的”。本文的動機(jī)正是在研究如何在知識蒸餾訓(xùn)練過程中避免或者說重復(fù)利用這種額外的計算結(jié)果，該文章的解決策略是提前保存每張圖片不同區(qū)域的軟監(jiān)督信號（regional soft label）在硬盤上，訓(xùn)練 student 過程同時讀取訓(xùn)練圖片和標(biāo)簽文件，從而達(dá)到復(fù)用標(biāo)簽的效果。所以問題就變成了：soft label 怎么來組織和存儲最為有效？下面具體來看該文章提出的策略。 1. FKD 算法框架介紹 FKD 框架的核心部分包含了兩個階段，如下圖：（1）軟標(biāo)簽（soft label）的生成和存儲；（2）使用軟標(biāo)簽（soft label）進(jìn)行模型訓(xùn)練。

如圖所示，上半部分展示了軟標(biāo)簽的生成過程，作者通過輸入多個 crops 進(jìn)入預(yù)訓(xùn)練好的 teacher 來產(chǎn)生需要的軟標(biāo)簽向量，同時作者還保存了：（1）每個 crop 對應(yīng)的坐標(biāo)和（2）是否翻轉(zhuǎn)的 Boolean 值。下半部分展示了 student 訓(xùn)練過程，作者在隨機(jī)采樣圖片的時候同時也會讀取它們對應(yīng)的軟標(biāo)簽文件，從中選取 N 個 crops 用于訓(xùn)練，額外數(shù)據(jù)增強(qiáng)比如 mixup，cutmix 會放在這個階段，從而節(jié)省了由于引入更多數(shù)據(jù)增強(qiáng)參數(shù)帶來的額外存儲開銷。 2. 采樣策略 本文還提出了一個 multi-crop sampling 的策略，即在一個 mini-batch 里面每張圖片采樣多個樣本 crops。當(dāng)總的訓(xùn)練 epochs 不變的前提下，該采樣方式可以大大減少數(shù)據(jù)讀取的次數(shù)，對于一些數(shù)據(jù)讀取不是非常高效或者產(chǎn)生嚴(yán)重瓶頸的集群設(shè)備，這種策略的加速效果非常明顯（如下表格所示）。同時在一張圖片采樣多個 crops 可以減少訓(xùn)練樣本間的方差，幫助穩(wěn)定訓(xùn)練，作者發(fā)現(xiàn)如果 crops 的數(shù)目不是太大的情況下可以明顯提升模型精度，但是一張圖片里面采樣太多 crops 數(shù)目會造成每個 mini-batch 里面訓(xùn)練樣本的信息差異不足（過于相似），因此過度采樣會影響性能，所以需要設(shè)置一個合理的數(shù)值。 3. 加速比 作者在實驗部分跟標(biāo)準(zhǔn)的訓(xùn)練方式以及 ReLabel 訓(xùn)練進(jìn)行了速度的比較，結(jié)果如下表格所示：可以看到，相比正常的分類框架，F(xiàn)KD 會快 16% 左右，而相比 ReLabel 則快了 30%，因為 ReLabel 相比正常訓(xùn)練需要讀取雙倍的文件數(shù)目。需要注意的是這個速度對比實驗中，F(xiàn)KD crop 數(shù)目為 4，如果選取更大的 crop 數(shù)目可以得到更高的加速比。

加速原因分析： 除了上述介紹的采用多個 crops 來進(jìn)行加速外，作者還分析了其他一些加速的因素，如下圖所示，ReLabel 在訓(xùn)練模型階段需要生成采樣數(shù)據(jù)的坐標(biāo)，同時需要使用 RoI-Align 和 Softmax 來生成所需的軟標(biāo)簽，相比而言，F(xiàn)KD 直接保存了坐標(biāo)信息和最終軟標(biāo)簽格式，因此讀取標(biāo)簽文件之后不需要做任何額外的后處理就可以直接訓(xùn)練，速度相比 ReLabel 也會更快。

4. 標(biāo)簽質(zhì)量分析 軟標(biāo)簽質(zhì)量是保證模型訓(xùn)練精度的一項最重要的指標(biāo)，作者通過可視化標(biāo)簽分布以及計算不同模型預(yù)測之間的交叉熵（cross-entropy）來證明了所提出的方式擁有更好的軟標(biāo)簽質(zhì)量。

上圖展示了 FKD 和 ReLabel 軟標(biāo)簽分布的情況對比，得到如下結(jié)論：

（第一行）FKD 相比 ReLabel 置信度更加平均也與輸入樣本內(nèi)容更加一致，作者分析原因是 ReLabel 將全局圖像輸入到模型中，而不是局部區(qū)域，這使得生成的全局標(biāo)簽映射編碼了更多全局類別信息同時忽略了背景信息，使得生成的軟標(biāo)簽過于接近單個語義標(biāo)簽。

（第二行）雖然存在一些樣本 ReLabel 和 FKD 之間的最大預(yù)測概率相似，但 FKD 包含更多標(biāo)簽分布中的從屬類別概率，而 ReLabel 的分布中并沒有捕獲這些從屬類別的信息。

（第三行）對于某些異常情況，F(xiàn)KD 比 ReLabel 更加健壯，例如目標(biāo)框含有松散邊界，或者只定位部分目標(biāo)等。

（第四行）在有些情況下，ReLabel 的標(biāo)簽分布意外的崩潰了（均勻分布），沒有產(chǎn)生一個主要的預(yù)測，而 FKD 仍然可以預(yù)測得很好。

5. 標(biāo)簽壓縮、量化策略 1）硬化 (Hardening)。在該策略中，樣本標(biāo)簽 Y_H 使用 teacher 預(yù)測的最大 logits 的索引。標(biāo)簽硬化策略產(chǎn)生的依然是 one-hot 的標(biāo)簽，如下公式所示：

2）平滑 (Smoothing)。平滑量化策略是將上述硬化后的標(biāo)簽 Y_H 替換為軟標(biāo)簽和均勻分布的分段函數(shù)組合，如下所示：

3）邊際平滑 (Marginal Smoothing with Top-K)。邊際平滑量化策略相比單一預(yù)測值保留了更多的邊際信息（Top-K）來平滑標(biāo)簽 Y_S：

4）邊際平滑歸一化 (Marginal Re-Norm with Top-K)。邊際平滑歸一化策略會將 Top-K 預(yù)測值重新歸一化到和為 1，并保持其他元素值為零（FKD 使用歸一化來校準(zhǔn) Top-K 預(yù)測值的和為 1，因為 FKD 存儲的軟標(biāo)簽是 softmax 處理之后的值）：

具體對應(yīng)上述各種量化策略的圖示如下圖所示：

6. 不同標(biāo)簽量化 / 壓縮策略的存儲大小比較 不同標(biāo)簽壓縮方法需要的存儲空間如下表格所示，所使用的數(shù)據(jù)集為 ImageNet-1K，其中 M 是軟標(biāo)簽生成階段每張圖像被采樣的數(shù)目，這里作者選取了 200 作為示例。Nim 是圖像數(shù)量， ImageNet-1K 數(shù)據(jù)集為 1.2M，SLM 是 ReLabel 標(biāo)簽矩陣的大小，Cclass 是類的數(shù)量，DDA 是需要存儲的數(shù)據(jù)增強(qiáng)的參數(shù)維度。

從表格中可以看到，在不做任何壓縮的情況下 FKD 軟標(biāo)簽需要的存儲空間為 0.9T，這在實際使用中顯然是不現(xiàn)實的，標(biāo)簽數(shù)據(jù)的大小已經(jīng)遠(yuǎn)遠(yuǎn)超過訓(xùn)練數(shù)據(jù)本身了。通過標(biāo)簽壓縮可以極大減少存儲大小，同時后面實驗也證明了合適的壓縮方式并不會損害模型精度。 7. 自監(jiān)督學(xué)習(xí)任務(wù)上的應(yīng)用 FKD 的訓(xùn)練方式也可以應(yīng)用于自監(jiān)督學(xué)習(xí)任務(wù)。作者使用自監(jiān)督算法比如 MoCo，SwAV 等來預(yù)訓(xùn)練 teacher 模型，然后按照上述方式生成用于自監(jiān)督的軟標(biāo)簽（unsupervised soft label），這個步驟跟監(jiān)督學(xué)習(xí)得到的 teacher 很相似。生成標(biāo)簽過程會保留原始自監(jiān)督模型中 projection head 并使用之后的最終輸出向量，然后將這個向量作為軟標(biāo)簽保存下來。得到該軟標(biāo)簽后，可以使用同樣的監(jiān)督式的訓(xùn)練方式來學(xué)習(xí)對應(yīng)的 student 模型。 8. 實驗結(jié)果 1）首先是在 ResNet-50 和 ResNet-101 上的結(jié)果，如下表所示，F(xiàn)KD 取得了 80.1%/ResNet-50 和 81.9%/ResNet-101 的精度。同時訓(xùn)練時間相比普通訓(xùn)練和 ReLabel 都快了很多。

2）作者還測試了 FKD 在 MEAL V2 上的結(jié)果，同樣得到了 80.91% 的結(jié)果。

3）Vision Transformer 上的結(jié)果： 接下來作者展示了在 vision transformer 上的結(jié)果，在不使用額外數(shù)據(jù)增強(qiáng)的情況下，F(xiàn)KD 就可以比之前知識蒸餾方法得到將近一個點的提升，同時訓(xùn)練速度快了 5 倍以上。

4）Tiny CNNs 上的結(jié)果：

5）消融實驗：

首先是不同壓縮策略，綜合考慮存儲需求和訓(xùn)練精度，邊際平滑策略是最佳的。

接下來是訓(xùn)練階段不同 crop 數(shù)目的對比，MEAL V2 由于使用了 pre-trained 的參數(shù)作為初始化權(quán)重，因此不同 crop 數(shù)目下性能都比較穩(wěn)定和接近。而 vanilla 和 FKD 在 crop=4 的時候表現(xiàn)得最好。尤其 vanilla，相比 crop=1 精度提升了一個點，crop 大于 8 之后精度下降明顯。

6）自監(jiān)督任務(wù)上的結(jié)果： 如下表所示，在自監(jiān)督學(xué)習(xí)任務(wù)上 FKD 方式還是可以很好的學(xué)習(xí)目標(biāo)模型，同時相比雙子結(jié)構(gòu)自監(jiān)督網(wǎng)絡(luò)訓(xùn)練和蒸餾訓(xùn)練，可以加速三到四倍。

9. 下游任務(wù) 下表給出了 FKD 模型在 ImageNet ReaL 和 ImageNetV2 兩個數(shù)據(jù)集上的結(jié)果，可以看到，F(xiàn)KD 在這些數(shù)據(jù)集上取得了穩(wěn)定的提升。

下表是 FKD 預(yù)訓(xùn)練模型在 COCO 目標(biāo)檢測任務(wù)上的結(jié)果，提升同樣明顯。

10. 可視化分析 如下兩張可視化圖所示，作者通過可視化中間特征層（attention map）的方式探索 FKD 這種 region-based 訓(xùn)練方式對模型產(chǎn)生的影響，作者對比了三種不同訓(xùn)練方式得到的模型：正常 one-hot label，ReLabel 和本文提出的 FKD。 （i） FKD 的預(yù)測的概率值相比 ReLabel 更加?。╯oft），因為 FKD 訓(xùn)練過程引入的上下文以及背景信息更多。在 FKD 隨機(jī) crop 的訓(xùn)練策略中，許多樣本采樣于背景（上下文）區(qū)域，來自 teacher 模型的軟預(yù)測標(biāo)簽更能真實的反映出實際輸入內(nèi)容，并且這些軟標(biāo)簽可能與 one-hot 標(biāo)簽完全不同，F(xiàn)KD 的訓(xùn)練機(jī)制可以更好的利用上下文中的額外信息。 （ii） FKD 的特征可視化圖在物體區(qū)域上具有更大的高響應(yīng)值區(qū)域，這表明 FKD 訓(xùn)練的模型利用了更多區(qū)域的線索進(jìn)行預(yù)測，進(jìn)而捕獲更多差異性和細(xì)粒度的信息。 （iii）ReLabel 的注意力可視化圖與 PyTorch 預(yù)訓(xùn)練模型更加接近，而 FKD 的結(jié)果跟他們相比具有交大差異性。這說明 FKD 方式學(xué)習(xí)到的注意力機(jī)制跟之前模型有著顯著的差別，從這點出發(fā)后續(xù)可以進(jìn)一步研究其有效的原因和工作機(jī)理。

審核編輯 ：李倩