論文解讀 Once for all: Train One Network and Specialize it for Efficient Deployment

1. 介紹

論文：https://arxiv.org/abs/1908.09791

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

早期的NAS算法大都是針對某一個具體的硬件平台做模型搜索。換句話說，對於新的平台則需要重新搜索，這樣既耗時又耗資源。所以這篇文章（簡稱OFA）的主要目的是只需要搜索一次就能為不同的平台找到合適的模型結構，所以論文名字取為 Once for All。

2. 搜索空間

OFA采用的也是權重共享的搜索策略，Supernet結構基於MobileNet的設計。

整個搜索空間（如Figure 2）主要包含三大塊：

• Depth (D)：Supernet可以看成是由多個Block組成的結構，而每個Block包含不同數量的layers，即為深度（Depth），論文中每個Block的layer數量搜索范圍是 [4,3,2]
• Width (W)：每個layer可以理解成是一個inverted residual module，其通道數由expand ratio控制，即為寬度（Width）,論文中每個layer的expand ratio搜索范圍是 [6,4,3]
• Kernel size (K)：這個很好理解，就是卷積核大小。論文中搜索范圍是 [7,5,3]
• OFA還額外搜索了不同分辨率的ImageNet數據輸入，這可不會對模型結構有什么影響，只不過能提高模型對不同大小的數據的泛化能力

OFA中的Supernet由5個block組成，而每個block的搜索空間大小是 \(S_i=(3\times3)^2+(3\times3)^3+(3\times3)^4\)，所以整個搜索空間大小為\(S=\prod_1^5S_i\approx2\times10^{19}\)

3. 搜索策略

再強調一遍，OFA的目的希望是只做一次搜索，之后就可以根據這個搜索的結果自動去找適合於指定平台的模型結構。

形式化的數學表達如下：

\[\min _{W_{o}} \sum_{a r c h_{i}} \mathcal{L}_{\text {val }}\left(C\left(W_{o}, \operatorname{arch}_{i}\right)\right) \]

說人話就是 讓所有候選架構在驗證集上的結果都盡可能的好，即最小化所有架構的期望損失函數之和。

但是搜索空間非常大，不可能遍歷所有模型，所以一個很naive的思路是是通過隨機采樣的方式來近似優化。但是如何采樣確實很有門道的，OFA就提出了 漸進式收縮 （progressive shrinking, PS）的采樣策略.

其實，PS 這個策略其實也很符合直覺，我們看Figure 2，它把整個訓練過程划分成多個phases：

1. Full Network：首先就是訓練最大的模型，即每個Block由4個layer組成，每個layer的通道expand ratio設置為6，kernel size設置為7。這個模型會作為Teacher net來做知識蒸餾。這一過程使用32個GPU訓練了180個epoch，batch size為2048，使用初始學習率是2.6的SGD優化器。
2. Elastic Kernel Size：之后固定D和W，搜索K，每次只采樣一個模型。訓練了125個epoch，初始學習率是0.96，其余參數和phase-1一樣
3. Elastic Depth：固定W,搜索D和K。這個phase又划分成了兩個stage：
   1. stage-1：D的搜索范圍是 [4, 3]，訓練了25個epoch，學習率是0.08
   2. stage-2： D的搜索范圍是 [4, 3, 2]，訓練了125個epoch，學習率是0.24
4. Elastic Width：同時搜索D,K,W。類似Phase-3，這個過程也分成了兩個stage，不再贅述。

4. 算法實現細節

4.1 搜索Kernel Size

卷積核大小搜索范圍是 [3,5,7]，在算法實現層面其實是共用一個7X7的卷積核，而要得到5X5的卷積核，OFA的做法是先截取中間的5X5的矩陣，然后把這個矩陣展開成一維向量，長度為25，然后再通過一個25X25的可學習的矩陣得到新的5X5的卷積核。3X3的卷積核的處理辦法類似，不過是基於變換后的5X5卷積再做變換得到的。

4.2 搜索Depth

上圖中的Unit可以理解成Supernet中的一個Block，所有Block初始化都是由4個layer組成。

在搜索depth的時候，假如搜索的結果是3，那么最后一個layer就不會參與計算。（如上圖第二個unit所示）

4.3 搜索Width

寬度的搜索其實就是通道數的搜索，OFA提出的策略是根據L1norm計算每個通道的重要性，然后對這些通道排序，最重要的排前面。比如總共有10個通道，某一時刻搜索的結果是只需要5個通道，那么就會對這個10個通道做排序，只用前5個通道做計算。我有試過這種策略，一旦對通道權重重新排序后，模型性能就會崩掉了，不過具體是我復現的代碼有問題還是其他問題就沒有細究了。

5. 實驗結果

微信公眾號：AutoML機器學習

MARSGGBO♥原創

如有意合作或學術討論歡迎私戳聯系~

郵箱:[e-mail]