【神經(jīng)網(wǎng)絡(luò)搜索】Once for all

【GiantPandaCV導(dǎo)語】Once for all是韓松組非常有影響力的工作，其最大的優(yōu)點是解耦了訓練和搜索過程，可以直接從超網(wǎng)中采樣出滿足一定資源限制的子網(wǎng)，而不需要重新訓練。該工作被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

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

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

也就是說，網(wǎng)絡(luò)設(shè)計主要受平臺的以下幾方面約束：

• 部署的硬件設(shè)備不同。
• 相同硬件設(shè)備上不同的部署要求，比如期望的延遲。

2. Contribution

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

3. Method

方法部分需要搞清楚兩個問題，一個是網(wǎng)絡(luò)是什么樣的？一個是網(wǎng)絡(luò)是如何訓練的？

第一個問題：once-for-all網(wǎng)絡(luò)長什么樣子？

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

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

由于網(wǎng)絡(luò)包括5個單元，所以候選的子網(wǎng)大概有個不同的子網(wǎng)，并且是在25個不同輸入分辨率下進行訓練。所有的子網(wǎng)都共享權(quán)重，只需要7.7M的參數(shù)量。

第二個問題：once-for-all怎樣才能同時訓練這么多子網(wǎng)絡(luò)？

由于once-for-all的目標是同時優(yōu)化所有的子網(wǎng)，所以需要考慮使用新的訓練策略。

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

通過以上分析可以看出，訓練超網(wǎng)是非常困難的，需要采用更好的訓練策略才能訓練得動超網(wǎng)。

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

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

以上策略的特點是：先訓練最大的，然后訓練小的。這樣可以盡可能減小訓練小模型的時候?qū)Υ竽Ｐ偷挠绊憽?/p>

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

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

通過以上描述可以看出來，權(quán)重共享的網(wǎng)絡(luò)優(yōu)化起來非常復(fù)雜，上邊的選擇的子網(wǎng)個數(shù)、學習率的選擇、epoch的選擇可能背后作者進行了無數(shù)次嘗試調(diào)參，才得到了一個比較好的結(jié)果。

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

簡單來說就是，中心共享+變換矩陣。從直覺上來講，優(yōu)化7x7的卷積核以后，再優(yōu)化中間的5x5卷積核勢必會影響原先7x7卷積核的結(jié)果，兩者在分布和數(shù)值上有較大的不同，強制訓練會導(dǎo)致性能有較大的下降，所以這就需要引入變換矩陣，具體實現(xiàn)是一個MLP，具體方法是：

• 不同層使用各自獨立的變換矩陣來共享權(quán)重。
• 相同層內(nèi)部，不同的通道之間共享變換矩陣。

2. Elastic Depth:  如何優(yōu)化不同深度的網(wǎng)絡(luò)呢？

深度為2、3、4的時候，按照上圖所示，選擇前i個層，進行訓練和優(yōu)化。

3. **Elastic Width: ** 如何優(yōu)化不同的通道個數(shù)？

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

部署階段的其他技術(shù)細節(jié)：

• 搜索子網(wǎng)，滿足一定的條件，比如延遲、功率等限制。
• 預(yù)測器：neural-network-twins, 功能是給定一個網(wǎng)絡(luò)結(jié)構(gòu)，預(yù)測其延遲和準確率。采樣了16K個不同架構(gòu)、不同分辨率的子網(wǎng)，然后再10K的驗證數(shù)據(jù)集上得到他們真實的準確率。【arch, accuracy】可以作為準確率預(yù)測器的訓練數(shù)據(jù)集。
• 構(gòu)建了一個延遲查找表 latency lookup table來預(yù)測不同目標硬件平臺的延遲。預(yù)測器訓練數(shù)據(jù)集只需要40GPU Days。

4. Experiment

訓練細節(jié)：

• 網(wǎng)絡(luò)搜索空間：MobileNetV3類似的
• 使用標準的SGD優(yōu)化器，momentum=0.9 weight decay=3e-5
• 初始學習率2.6 使用cosine schedule來進行l(wèi)earning rate decay
• 在32GPU上使用2048的batch size訓練了180個epoch
• 在V100GPU上訓練了1200個GPU hours

漸進收縮策略：

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

實驗結(jié)果：

once-for-all在ImageNet上訓練結(jié)果如上表所示，其中可以發(fā)現(xiàn)雖然訓練代價比較高，但是搜索的代價穩(wěn)定在40GPU Hours并且取得了相同量級MACs下不錯結(jié)果。

5. Revisiting

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

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%20Slides.pdf

對文章有疑問或者研究NAS的同學可以添加筆者微信