【神經網絡架構搜索】ProxylessNAS直接在ImageNet上搜索

【GiantPandaCV導語】這篇是MIT韓松實驗室發佈的文章，是第一個直接在ImageNet上進行搜索的NAS算法，並且提出了直接在目標硬件上對latency進行優化的方法。相比於同期算法NASNet、MnasNet等，搜索代價降低了200倍。

0. Info

Title: ProxylessNAS: Direct Neural Architecture Search On Target Task and Hardware

Author: MIT韓松團隊

Link: https://arxiv.org/pdf/1812.00332v2

Date: ICLR2019

Code: https://github.com/MIT-HAN-LAB/ProxylessNAS

1. Motivation

之前的算法往往是在Proxy任務(比如CIFAR10、使用更少的epoch等方案)上進行訓練，然後遷移到目標數據集(比如ImageNet)，但是這種方案是次優的，存在一定的gap。

並且之前的Gradient-based NAS比如Darts由於需要將整個超網放在顯存中，對資源的佔用會比較大，因此對網絡的大小帶來了限制。

爲了解決以上痛點，提出了ProxylessNAS，可以直接爲目標硬件平臺優化網絡，解決了Gradient-based NAS中存在的高內存消耗問題，降低了計算代價。

2. Contribution

ProxylessNAS將搜索代價從40000GPU hours壓縮到200GPU hours。

文章主要共享有：

ProxylessNAS是第一個不借助proxy任務，直接在目標數據集ImageNet上進行學習的NAS算法。可以有效地擴大搜索空間，達到更好的表現。
提出了從Path-level Pruning角度來處理NAS的思路，解釋了NAS和模型壓縮之間存在的緊密聯繫。通過使用Path-level binarization節省了一個數量級的內存消耗。
提出了Gradient-based方法來處理硬件目標，比如latency。讓Latency這種硬件指標也變得可微分。並且ProxylessNAS也是第一個研究不同硬件架構上專用神經網絡架構的工作。
使用一系列實驗來證明ProxylessNAS的直接性directness和定製性specialization。在ImageNet和CIFAR10數據集上取得了SOTA的結果。

3. Method

3.1 過參數網絡構建

參考了one-shot和darts的搜索空間，構建了過參數網絡，

\[ m_ {O}^ {One-Shot} (x)= \sum _ {i=1}^ {n} O_ {i} (x) \]

對於one-shot來說，會將所有path路徑輸出相加；

\[ m_ {O}^ {DARTs} (x)= \sum _ {i=1}^ {N} p_ {i} o_ {i} (x)= \sum _ {i=1}^ {N} \frac {exp(\alpha _ {i})}{\sum_j exp(\alpha_j) } O_{i} (x). \]

對於darts來說，使用的是weighted sum的方式得到輸出。

以上兩種方式的實現都需要將中間結果保存，即所有路徑的結果，這樣相當於要佔用N倍的顯存。

爲了解決以上問題，提出了Path binarization。

3.2 Path-level binarization

爲了降低訓練過程中佔用的內存，每次訓練過參數網絡中的單條路徑分支。

Path binarization需要引入N個網絡架構參數α，pi的計算和darts一直。除此以外還有二進制開關g，1代表保留該路徑，0代表不啓用該路徑。

然後得到輸出可以看出，在運行過程中，只有單路徑被激活，節省了一個量級的顯存消耗。

3.3 訓練二值化網絡

和darts類似，ProxylessNAS交替更新權重參數和架構參數。

在訓練權重參數的時候，固定住架構參數，然後隨機採樣二進制開關。然後激活路徑對應的網絡使用梯度下降算法在訓練集上進行訓練。
在訓練架構參數的時候，固定住權重參數，然後更新架構參數。

完成訓練以後，選擇擁有最高path weight的path作爲最終網絡結構。

在計算梯度的時候，進行了以下近似：

\[\frac{\partial L}{\partial \alpha_{i}}=\sum_{j=1}^{N} \frac{\partial L}{\partial p_{j}} \frac{\partial p_{j}}{\partial \alpha_{i}} \approx \sum_{j=1}^{N} \frac{\partial L}{\partial g_{j}} \frac{\partial p_{j}}{\partial \alpha_{i}}=\sum_{j=1}^{N} \frac{\partial L}{\partial g_{j}} \frac{\partial\left(\frac{\exp \left(\alpha_{j}\right)}{\sum_{k} \exp \left(\alpha_{k}\right)}\right)}{\partial \alpha_{i}}=\sum_{j=1}^{N} \frac{\partial L}{\partial g_{j}} p_{j}\left(\delta_{i j}-p_{i}\right) \]

其中\(\delta_{i j}\)是一個指示器，當i=j的時候爲1，反之爲0。

但是如果用以上公式進行更新，實際上還是需要N背的計算資源。所以這裏每次更新架構參數的時候，採樣兩條路徑，相當於將N降到了2。

ps: 感覺這種方式會導致收斂變慢，相比於Darts，爲了內存空間犧牲了時間效率。

3.4 使硬件指標Latency可微分

提出兩個方法來解決不可微的目標。

方法一：讓Latency可微

F代表對應op的延遲預測器，預測op所對應的延遲。

\[\mathbb{E}\left[\text { latency }_{i}\right]=\sum_{j} p_{j}^{i} \times F\left(o_{j}^{i}\right) \]

求偏導：

\[\partial \mathbb{E}\left[\text { latency }_{i}\right] / \partial p_{j}^{i}=F\left(o_{j}^{i}\right) \]

將Latency作爲loss一項：

\[\text { Loss }=\operatorname{Loss}_{C E}+\lambda_{1}\|w\|_{2}^{2}+\lambda_{2} \mathbb{E}[\text { latency }] \]

方法二：使用基於強化學習方法

使用REINFORCE算法來訓練二值化權重。

\[\begin{aligned} J(\alpha) &=\mathbb{E}_{g \sim \alpha}\left[R\left(\mathcal{N}_{g}\right)\right]=\sum_{i} p_{i} R\left(\mathcal{N}\left(e=o_{i}\right)\right) \\ \nabla_{\alpha} J(\alpha) &=\sum_{i} R\left(\mathcal{N}\left(e=o_{i}\right)\right) \nabla_{\alpha} p_{i}=\sum_{i} R\left(\mathcal{N}\left(e=o_{i}\right)\right) p_{i} \nabla_{\alpha} \log \left(p_{i}\right) \\ &=\mathbb{E}_{g \sim \alpha}\left[R\left(\mathcal{N}_{g}\right) \nabla_{\alpha} \log (p(g))\right] \approx \frac{1}{M} \sum_{i=1}^{M} R\left(\mathcal{N}_{g^{i}}\right) \nabla_{\alpha} \log \left(p\left(g^{i}\right)\right) \end{aligned} \]

4. Experiment

CIFAR10: 選擇PyramidNet作爲backbone，具體實現細節需要看原文。

R代表使用強化學習得到的結果，G代表使用梯度優化得到的結果。c/o代表使用了cutout。

ImageNet: 在GPU、Mobile phone、CPU上進行了實驗。

主要關注Search Cost這一欄，相比之前的MNasNet、NASNte、AmoebaNet節省了200倍的GPU hours。

上圖結果展示了ProxylessNAS的優越性，其訓練的時間和佔用的內存都和正常訓練屬於一個量級。

ProxylessNAS專門爲不同硬件設計的網絡在對應硬件上的延遲是最低的。

5. Revisiting

ProxylessNAS是第一個直接在目標數據集ImageNet上進行訓練的神經網絡搜索算法，通過使用path binarization解決了以往Gradient-based NAS中存在的顯存佔用過高的問題,可以做到佔用和普通網絡相同量級顯存。同時ProxylessNAS也是第一個可以爲不同的硬件設置定製化網絡的算法。

在這個部分補充一些知識點：深度學習中的顯存佔用分析。並回答幾個問題。

問題一：哪些操作會佔用顯存？

模型自身的參數，主要包括權重。有參數的層：卷積、全連接、BN等都包含了大量參數。這部分內容和輸入設置無關，模型加載後就會佔用顯存。
模型中間的結果、梯度等。優化器在優化過程中需要保存對應權重的梯度纔可以計算，不同優化器需要的中間變量數目不定，比如Adam優化器，動量佔用顯存最多。
輸入輸出佔用的顯存。主要包括feature map等，計算得到每一層輸出的tensor形狀，就可以計算相應需要的顯存佔用。這部分顯存佔用是和batch size成正比的。

問題二：顯存大小和batch size是否成正比？

有了上邊問題做鋪墊，這個問題就很好解釋了。顯存大小包括三個部分，模型自身參數、模型梯度、輸入輸出佔用顯存。batch size僅僅會影響到輸入輸出的顯存佔用，所以並不是完全成正比。

問題三：如何降低顯存佔用？

經常會遇到GPU顯存溢出out of memory的報錯，通過以上分析，可以得知有以下處理方法：

使用更小的模型，可以修改模型通道數或者更換模型。
降低batch size大小。
可以使用amp中的混合精度方法，通過使用fp16降低顯存佔用。
清空中間變量，優化代碼比如inplace=True

問題四：增大batch size會帶來什麼影響？

batch size增大可以加快速度，但是是並行計算方面的，帶來的增益優先。
batch size增大可以減緩梯度震盪，需要更少的迭代次數，收斂更快。但是每次迭代耗時更長。
batch size調整的同時也需要調整對應的learning rate，以ImageNet上ResNet50爲例，官方推薦的lr=0.1 對照 batch size=256, 假設我要調整batch size=1024，那麼新的lr=0.1 x 1024 / 256。