OREPA:阿里提出訓練也很快的重參數策略,內存減半,速度加倍 | CVPR 2022

論文提出了在線重參數方法OREPA,在訓練階段就能將複雜的結構重參數爲單卷積層,從而降低大量訓練的耗時。爲了實現這一目標,論文用線性縮放層代替了訓練時的BN層,保持了優化方向的多樣性和特徵表達能力。從實驗結果來看,OREPA在各種任務上的準確率和效率都很不錯

來源:曉飛的算法工程筆記 公衆號

論文: Online Convolutional Re-parameterization

Introduction


  除了準確率外,模型的推理速度也是很重要的。爲了獲得部署友好且精度高的模型,近期很多研究提出基於結構重參數化來提高模型性能。用於結構重參數化的模型在訓練階段和推理階段具有不同的結構,訓練時使用複雜的結構來獲得高精度,而訓練後通過等效變換將一個複雜的結構壓縮成能夠快速推理的線性層。壓縮後的模型通常具有簡潔的架構,例如類似VGG的或類似ResNet的結構。從這個角度來看,重參數化策略可以在不引入額外推理時間成本的情況下提高模型性能。公衆號之前發了RepVGG的論文解讀《RepVGG:VGG,永遠的神! | 2021新文》,有興趣可以去看看。

  BN層是重參數模型中的關鍵組成部分,在每個卷積層之後添加一個BN層,如果圖1b所示,移除BN層會導致嚴重的精度下降。在推理階段,複雜的結構可以被壓縮到單個卷積層中。而在訓練階段,由於BN層需要非線性地將特徵圖除以其標準差,只能單獨計算每個分支。因此,存在大量中間計算操作(大FLOPS)和緩衝特徵圖(高內存使用),帶來巨大的計算開銷。更糟糕的是,高額的訓練消耗阻礙了探索更復雜和可能更強大的重新參數結構。
  爲什麼BN層對重參數化如此重要?根據實驗和分析,論文發現BN層中的縮放因子能夠使不同分支的優化方向多樣化。基於這個發現,論文提出了在線重參數化方法OREPA,如圖1c所示,包含兩個步驟:

  • block linearization:去掉所有非線性泛化層,轉而引入線性縮放層。線性縮放層不僅能與BN層一樣使不同分支的優化方向多樣化,還可以在訓練時合併計算。
  • block squeezing:將複雜的線性結構簡化爲單個卷積層。

  OREPA減少了由中間層引起的計算和存儲開銷,能夠顯著降低訓練消耗(65%-75%顯存節省、加速1.5-2.3倍)且對性能的影響很小,使得探索更復雜的重參數化結果成爲可能。爲了驗證這一點,論文進一步提出了幾個重新參數化的組件以獲得更好的性能。
  論文的貢獻包含以下三點:

  • 提出在線重參數化方法OREPA,能夠大幅提高重參數化模型的訓練效率,使得探索更強的重參數結構成爲可能。
  • 根據對重參數模型原理的分析,將BN層替換爲線性縮放層,保持優化方向多樣化特性和特徵表達能力。
  • 通過各種視覺任務實驗表明,OREPA在準確率和訓練效率方面都優於以前的重參數化模型。

Online Re-Parameterization



  OREPA能夠將訓練期間的複雜結構簡化爲單一卷積層,維持準確率不變。OREPA的變換流程如圖2所示,包含block linearization和block squeezing兩個步驟。

Preliminaries: Normalization in Re-param

  BN層是重參數中多層和多分支結構的關鍵結構,是重參數模型性能的基礎。以DBB和RepVGG爲例,去掉BN層後(改爲多分支後統一進行BN操作)性能會有明顯的下降,如表1所示。
  比較意外的是,BN層的使用會帶來過高的訓練消耗。在推理階段,重參數結構中的所有中間操作都是線性的,可以進行合併計算。而在訓練階段,由於BN層是非線性的(需要除以特徵圖的標準差),無法進行合併計算。無法合併就會導致中間操作需要單獨計算,產生巨大的計算消耗和內存成本。而且,過高的成本也阻礙了更復雜的結構的探索。

Block Linearization

  雖然BN層阻止了訓練期間的合併計算,但由於準確率問題,仍然不能直接將其刪除。爲了解決這個問題,論文引入了channel-wise的線性縮放作爲BN層的線性替換,通過可學習的向量進行特徵圖的縮放。線性縮放層具有BN層的類似效果,引導多分支向不同方向進行優化,這是重參數化性能的核心。

  基於線性縮放層,對重參數化結構進行修改,如圖3所示,以下三個步驟:

  • 移除所有非線性層,即重參數化結構中的歸一化層。
  • 爲了保持優化的多樣性,在每個分支的末尾添加了一個縮放層,即BN層的線性替代。
  • 爲了穩定訓練過程,在所有分支之後添加一個BN層。

  經過block linearization操作後,重參數結構中就只存在線性層,這意味着可以在訓練階段合併結構中的所有組件。

Block Squeezing

  Block squeezing將計算和內存過多的中間特徵圖上的操作轉換爲更快捷的單個卷積核核操作,這意味着在計算和內存方面將重參數的額外訓練成本從\(O(H\times W)\)減少到\(O(KH\times KW )\),其中\((KH, KW)\)是卷積核的形狀。
  一般來說,無論線性重參數結構多複雜,以下兩個屬性都始終成立:

  • 重參數結構中的所有線性層(例如深度卷積、平均池化和建議的線性縮放)都可以用具有相應參數的卷積層來表示,具體證明可以看原文的附錄。
  • 重參數結構可表示爲一組並行分支,每個分支包含一串卷積層。

  有了上述兩個屬性,就以將多層(即順序結構)和多分支(即並行結構)壓縮爲單個卷積,如圖4a和圖4b所示。原文有部分轉換的公式證明,有興趣的可以去看看原文對應章節,這塊不影響對Block Squeezing的思想的理解。

Gradient Analysis on Multi-branch Topology

  論文從梯度回傳的角度對多分支與block linearization的作用進行了分析,裏面包含了部分公式推導,有興趣的可以去看看原文對應章節。這裏總結主要的兩個結論:

  • 如果使用分支共享的block linearization,多分支的優化方向和幅度與單分支一樣。
  • 如果使用分支獨立的block linearization,多分支的優化方向和幅度與單分支不同。

  上面的結論表明了block linearization步驟的重要性。當去掉BN層後,縮放層能夠保持優化方向的多樣化,避免多分支退化爲單分支。

Block Design

  由於OREPA節省了大量訓練消耗,爲探索更復雜的訓練結構提供了可能性。論文基於DBB設計了全新的重參數模塊OREPA-ResNet,加入了以下組件:

  • Frequency prior filter:Fcanet指出池化層是頻域濾波的一個特例,參考此工作加入1x1卷積+頻域濾波分支。
  • Linear depthwise separable convolution:對深度可分離卷積進行少量修改,去掉中間的非線性激活以便在訓練期間合併。
  • Re-parameterization for 1x1 convolution:之前的研究主要關注3×3卷積層的重參數而忽略了1×1卷積,但1x1卷積在bottleneck結構中十分重要。其次,論文添加了一個額外的1x1卷積+1x1卷積分支,對1x1卷積也進行重參數。
  • Linear deep stem:一般網絡採用7x7卷積+3x3卷積作爲stem,有的網絡將其替換爲堆疊的3個3x3卷積取得了不錯的準確率。但論文認爲這樣的堆疊設計在開頭的高分辨率特徵圖上的計算消耗非常高,爲此將3個3x3卷積與論文提出的線性層一起壓縮爲單個7x7卷積層,能夠大幅降低計算消耗並保存準確率。

  OREPA-ResNet中的block設計如圖6所示,這應該是一個下采樣的block,最終被合併成單個3x3卷積進行訓練和推理。

Experiment


  各組件對比實驗。

  縮放層對各層各分支的相似性的影響。

  線性縮放策略對比,channel-wise的縮放最好。

  在線和離線重參數的訓練耗時對比。

  與其他重參數策略進行對比。

  在檢測和分割任務上進行對比。

Conclusion


  論文提出了在線重參數方法OREPA,在訓練階段就能將複雜的結構重參數爲單卷積層,從而降低大量訓練的耗時。爲了實現這一目標,論文用線性縮放層代替了訓練時的BN層,保持了優化方向的多樣性和特徵表達能力。從實驗結果來看,OREPA在各種任務上的準確率和效率都很不錯。



如果本文對你有幫助,麻煩點個贊或在看唄~
更多內容請關注 微信公衆號【曉飛的算法工程筆記】

work-life balance.

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章