首個(gè)快速知識蒸餾的視覺框架：ResNet50 80.1%精度，訓(xùn)練加速30%

• 論文和項(xiàng)目****：http://zhiqiangshen.com/projects/FKD/index.html
• 代碼：https://github.com/szq0214/FKD

計(jì)算機(jī)視覺研究院專欄

作者：Edison_G

今天介紹一篇來自卡耐基梅隆大學(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)的知識蒸餾策略之一，代碼和模型已全部開源！

轉(zhuǎn)自《機(jī)器之心》

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

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

加速原因分析：除了上述介紹的采用多個(gè) crops 來進(jìn)行加速外，作者還分析了其他一些加速的因素，如下圖所示，ReLabel 在訓(xùn)練模型階段需要生成采樣數(shù)據(jù)的坐標(biāo)，同時(shí)需要使用 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)練精度的一項(xiàng)最重要的指標(biāo)，作者通過可視化標(biāo)簽分布以及計(jì)算不同模型預(yù)測之間的交叉熵（cross-entropy）來證明了所提出的方式擁有更好的軟標(biāo)簽質(zhì)量。

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

• （第一行）FKD 相比 ReLabel 置信度更加平均也與輸入樣本內(nèi)容更加一致，作者分析原因是 ReLabel 將全局圖像輸入到模型中，而不是局部區(qū)域，這使得生成的全局標(biāo)簽映射編碼了更多全局類別信息同時(shí)忽略了背景信息，使得生成的軟標(biāo)簽過于接近單個(gè)語義標(biāo)簽。
• （第二行）雖然存在一些樣本 ReLabel 和 FKD 之間的最大預(yù)測概率相似，但 FKD 包含更多標(biāo)簽分布中的從屬類別概率，而 ReLabel 的分布中并沒有捕獲這些從屬類別的信息。
• （第三行）對于某些異常情況，F(xiàn)KD 比 ReLabel 更加健壯，例如目標(biāo)框含有松散邊界，或者只定位部分目標(biāo)等。
• （第四行）在有些情況下，ReLabel 的標(biāo)簽分布意外的崩潰了（均勻分布），沒有產(chǎn)生一個(gè)主要的預(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，因?yàn)?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，這在實(shí)際使用中顯然是不現(xiàn)實(shí)的，標(biāo)簽數(shù)據(jù)的大小已經(jīng)遠(yuǎn)遠(yuǎn)超過訓(xùn)練數(shù)據(jù)本身了。通過標(biāo)簽壓縮可以極大減少存儲大小，同時(shí)后面實(shí)驗(yàn)也證明了合適的壓縮方式并不會損害模型精度。
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），這個(gè)步驟跟監(jiān)督學(xué)習(xí)得到的 teacher 很相似。生成標(biāo)簽過程會保留原始自監(jiān)督模型中 projection head 并使用之后的最終輸出向量，然后將這個(gè)向量作為軟標(biāo)簽保存下來。得到該軟標(biāo)簽后，可以使用同樣的監(jiān)督式的訓(xùn)練方式來學(xué)習(xí)對應(yīng)的 student 模型。
8. 實(shí)驗(yàn)結(jié)果
1）首先是在 ResNet-50 和 ResNet-101 上的結(jié)果，如下表所示，F(xiàn)KD 取得了 80.1%/ResNet-50 和 81.9%/ResNet-101 的精度。同時(shí)訓(xùn)練時(shí)間相比普通訓(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 就可以比之前知識蒸餾方法得到將近一個(gè)點(diǎn)的提升，同時(shí)訓(xùn)練速度快了 5 倍以上。

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

5）消融實(shí)驗(yàn)：首先是不同壓縮策略，綜合考慮存儲需求和訓(xùn)練精度，邊際平滑策略是最佳的。接下來是訓(xùn)練階段不同 crop 數(shù)目的對比，MEAL V2 由于使用了 pre-trained 的參數(shù)作為初始化權(quán)重，因此不同 crop 數(shù)目下性能都比較穩(wěn)定和接近。而 vanilla 和 FKD 在 crop=4 的時(shí)候表現(xiàn)得最好。尤其 vanilla，相比 crop=1 精度提升了一個(gè)點(diǎn)，crop 大于 8 之后精度下降明顯。

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

9. 下游任務(wù)
下表給出了 FKD 模型在 ImageNet ReaL 和 ImageNetV2 兩個(gè)數(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 更加小（soft），因?yàn)?FKD 訓(xùn)練過程引入的上下文以及背景信息更多。在 FKD 隨機(jī) crop 的訓(xùn)練策略中，許多樣本采樣于背景（上下文）區(qū)域，來自 teacher 模型的軟預(yù)測標(biāo)簽更能真實(shí)的反映出實(shí)際輸入內(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ī)制跟之前模型有著顯著的差別，從這點(diǎn)出發(fā)后續(xù)可以進(jìn)一步研究其有效的原因和工作機(jī)理。