【神經網絡搜索】Once for all

【GiantPandaCV導語】Once for all是韓松組非常有影響力的工作,其最大的優點是解耦了訓練和搜索過程,可以直接從超網中採樣出滿足一定資源限制的子網,而不需要重新訓練。該工作被ICLR20接收。

0. Info

Title: Once-for-All: Train one Network and Specialize it for Efficient Deployment

Author: 韓松組

Link: https://arxiv.org/pdf/1908.09791v5

Publish: ICLR2020

1. Motivation

傳統網絡搜索方法往往只能針對某個特定設備或者特定資源限制的平臺進行鍼對性搜索。對於不同的設備,往往需要在該設備上從頭訓練。這樣的方法擴展性很差並且計算代價太大,所以once for all從這個角度出發,希望能做到將訓練和搜索過程解耦,從而可以訓練一個支持不同架構配置的once-for-all網絡(類似超網的概念),通過從once-for-all網絡中選擇一個子網,就可以在不需要額外訓練的情況下得到一個專門的子網絡。

  • 不同的硬件平臺有着不同的硬件效率限制,比如延遲、功耗
  • 不同的硬件平臺的硬件資源差別很大,比如最新的手機和最老的手機。
  • 相同的硬件上,不同的電池條件、工作負載下使用的網絡模型也是不同的。

也就是說,網絡設計主要受平臺的以下幾方面約束:

  • 部署的硬件設備不同。
  • 相同硬件設備上不同的部署要求,比如期望的延遲。

2. Contribution

  • 提出瞭解決以上問題的方法:設計once-for-all網絡,可以在不同的網絡配置下進行部署。
  • 推理過程使用的模型是once-for-all網絡的一部分,可以無需重新訓練就能靈活的支持深度、寬度、卷積核大小、分辨率等參數的不同。
  • 提出了漸進式收縮的訓練策略來訓練once-for-all網絡

3. Method

方法部分需要搞清楚兩個問題,一個是網絡是什麼樣的?一個是網絡是如何訓練的?

第一個問題:once-for-all網絡長什麼樣子?

once-for-all網絡支持深度、寬度、卷積核大小、圖像分辨率四個因素的變化。

  • elastic depth: 代表選擇任意深度的網絡,每個單元的深度有{2,3,4}三個選項。
  • elastic width: 代表選擇任意數量的通道,寬度比例有{3,4,6}三個選項。
  • elastic kernel size: 代表選擇任意的卷積核大小,有{3,5,7}三個選項。
  • arbitrary resolution: 代表圖像的分辨率是可變的,從128到224,stride=4的分辨率均可。

由於網絡包括5個單元,所以候選的子網大概有\(((3\times 3)^2+(3\times 3)^3+(3\times 3)^4)^5 = 2\times 10^{19}\)個不同的子網,並且是在25個不同輸入分辨率下進行訓練。所有的子網都共享權重,只需要7.7M的參數量。

第二個問題:once-for-all怎樣才能同時訓練這麼多子網絡?

由於once-for-all的目標是同時優化所有的子網,所以需要考慮使用新的訓練策略。

  • 最簡單的想法:不考慮計算代價的情況下,每次梯度的更新都是由全體子網計算得到的。雖然這樣最準確,但是可想而知計算代價過高,並不實際。
  • 可行的想法:每次梯度是由一部分子網計算得到的。筆者曾經嘗試過這種方法(single path one shot),收斂的速度非常慢,得到的準確率也非常低。這很可能是在訓練過程中,由於權重是共享的,梯度在同一個參數的更新上可能帶來衝突,減緩了訓練的過程,並且達到最終的準確率也不夠高。

通過以上分析可以看出,訓練超網是非常困難的,需要採用更好的訓練策略才能訓練得動超網。

本文提出了Progressive Shrinking策略來解決以上問題,如下圖所示:

  1. 先訓練最大的kernel size, depth , width的網絡
  2. 微調網絡來支持子網,即將小型的子網加入採樣空間中。比如說,當前正在微調kernel size的時候,其他的幾個選項depth, width需要維持最大的值。另外,分辨率大小是每個batch隨機採樣的,類似於yolov3裏的訓練方法。
  3. 採用了知識蒸餾的方法,讓最大的超網來指導子網的學習。

漸進式收縮訓練策略

以上策略的特點是:先訓練最大的,然後訓練小的。這樣可以儘可能減小訓練小模型的時候對大模型的影響。

下面對照上圖詳細展開PS策略:

  • 訓練整個網絡,最大kernel,最寬channel,最深depth
  • 訓練可變kernel size, 每次採樣一個子網,使用0.96的初始學習率訓練125個epoch
  • 訓練可變depth,採樣兩個子網,每次更新收集兩者的梯度。第一個stage使用0.08的學習率訓練25個epoch;使用0.24的學習率訓練125個epoch。
  • 訓練可變width,採樣四個子網,每次更新收集四個子網梯度。第一個stage使用0.08的學習率訓練25個epoch;使用0.24的學習率訓練125個epoch。

通過以上描述可以看出來,權重共享的網絡優化起來非常複雜,上邊的選擇的子網個數、學習率的選擇、epoch的選擇可能背後作者進行了無數次嘗試調參,纔得到了一個比較好的結果。

  1. **Elastic Kernel Size: **Kernel Size是如何共享的呢?

卷積核共享

簡單來說就是,中心共享+變換矩陣。從直覺上來講,優化7x7的卷積核以後,再優化中間的5x5卷積核勢必會影響原先7x7卷積核的結果,兩者在分佈和數值上有較大的不同,強制訓練會導致性能有較大的下降,所以這就需要引入變換矩陣,具體實現是一個MLP,具體方法是:

  • 不同層使用各自獨立的變換矩陣來共享權重。
  • 相同層內部,不同的通道之間共享變換矩陣。
  1. Elastic Depth: 如何優化不同深度的網絡呢?

不同的深度優化

深度爲2、3、4的時候,按照上圖所示,選擇前i個層,進行訓練和優化。

  1. **Elastic Width: ** 如何優化不同的通道個數?

通道的選擇

並沒有使用類似slimmable network中那種選取前n個通道的策略,而是選取了一個channel importance進行排序,通道重要性計算方法是L1範數,L1範數越大,代表其重要性比較高,選擇重要性最高的前n個通道。

部署階段的其他技術細節:

  • 搜索子網,滿足一定的條件,比如延遲、功率等限制。
  • 預測器:neural-network-twins, 功能是給定一個網絡結構,預測其延遲和準確率。
    採樣了16K個不同架構、不同分辨率的子網,然後再10K的驗證數據集上得到他們真實的準確率。【arch, accuracy】可以作爲準確率預測器的訓練數據集。
  • 構建了一個延遲查找表 latency lookup table來預測不同目標硬件平臺的延遲。
    預測器訓練數據集只需要40GPU Days。

4. Experiment

訓練細節

  • 網絡搜索空間:MobileNetV3類似的
  • 使用標準的SGD優化器,momentum=0.9 weight decay=3e-5
  • 初始學習率2.6 使用cosine schedule來進行learning rate decay
  • 在32GPU上使用2048的batch size訓練了180個epoch
  • 在V100GPU上訓練了1200個GPU hours

漸進收縮策略:

漸進收縮策略的有效性

上圖展示了使用漸進收縮策略以後帶來的性能提升,可以看出,不同的架構配置下,都帶來了2-4%的性能提升。

實驗結果:

實驗結果

once-for-all在ImageNet上訓練結果如上表所示,其中可以發現雖然訓練代價比較高,但是搜索的代價穩定在40GPU Hours並且取得了相同量級MACs下不錯結果。

5. Revisiting

  • 並沒有什麼突破性的創新點,但是每個點都做得很紮實,在一個不錯的motivation下,將故事講的非常引人入勝。所以會講故事+紮實的實驗結果+(大量的算力) 纔得到這個結果。
  • 漸進收縮策略中先訓練kernel size,在訓練depth,最後訓練width的順序並沒有明確指出爲何是這樣的順序。
  • 通道的搜索策略筆者把它搬到single path one shot上進行了實驗,效果並不理想。
  • 通道搜索策略中once for all計算L1 Norm是根據輸入的通道來計算的,有點違背直覺,通常來講根據輸出通道計算更符合直覺一些。這一點可以參考通道剪枝,可能兩種方法都是可行的,具體選哪個需要看實驗結果。
  • 共享kernel size那部分工作的分析非常好,想到使用一個轉移矩陣來適應不同kernel所需要的分佈非常符合直覺。
  • 這篇工作代碼量非常大,非常的工程化,從文章的實驗也能看出裏邊需要非常強的工程能力,調參能力、才能在頂會上發表。

6. Reference

https://zhuanlan.zhihu.com/p/164695166

https://github.com/mit-han-lab/once-for-all

https://arxiv.org/abs/1908.09791

https://file.lzhu.me/projects/OnceForAll/OFA Slides.pdf

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