麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

機器之心發表於2018-12-07
  • 論文:ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware.
  • 網站:https://hanlab.mit.edu/projects/proxylessNAS/

  • 論文: https://arxiv.org/pdf/1812.00332.pdf

  • 程式碼:https://github.com/MIT-HAN-LAB/ProxylessNAS

摘要:NAS 受限於其過高的計算資源 (GPU 時間, GPU 記憶體) 需求,仍然無法在大規模任務 (例如 ImageNet) 上直接進行神經網路結構學習。目前一個普遍的做法是在一個小型的 Proxy 任務上進行網路結構的學習,然後再遷移到目標任務上。這樣的 Proxy 包括: (i) 訓練極少量輪數; (ii) 在較小的網路下學習一個結構單元 (block),然後透過重複堆疊同樣的 block 構建一個大的網路; (iii) 在小資料集 (例如 CIFAR) 上進行搜尋。然而,這些在 Proxy 上最佳化的網路結構在目標任務上並不是最優的。在本文中,我們提出了 ProxylessNAS,第一個在沒有任何 Proxy 的情況下直接在 ImageNet 量級的大規模資料集上搜尋大設計空間的的 NAS 演算法,並首次專門為硬體定製 CNN 架構。我們將模型壓縮 (減枝,量化) 的思想與 NAS 進行結合,把 NAS 的計算成本 (GPU 時間, GPU 記憶體) 降低到與常規訓練相同規模,同時保留了豐富的搜尋空間,並將神經網路結構的硬體效能 (延時,能耗) 也直接納入到最佳化目標中。我們在 CIFAR-10 和 ImageNet 的實驗驗證了」直接搜尋」和「為硬體定製」的有效性。在 CIFAR-10 上,我們的模型僅用 5.7M 引數就達到了 2.08% 的測試誤差。對比之前的最優模型 AmoebaNet-B,ProxylessNAS 僅用了六分之一的引數量就達到了更好的結果。在 ImageNet 上,ProxylessNAS 比 MobilenetV2 高了 3.1% 的 Top-1 正確率,並且在 GPU 上比 MobilenetV2 快了 20%。在同等的 top-1 準確率下 (74.5% 以上), ProxylessNAS 的手機實測速度是當今業界標準 MobileNetV2 的 1.8 倍。在用 ProxylessNAS 來為不同硬體定製神經網路結構的同時,我們發現各個平臺上搜尋到的神經網路在結構上有很大不同。這些發現為之後設計高效 CNN 結構提供新的思路。

背景

近期研究中,神經結構搜尋(NAS)已經在各種深度學習任務(例如影像識別)的神經網路結構設計自動化方面取得了很大成功。然而,傳統 NAS 演算法的計算量需求往往令人望而卻步,例如 NASNet 需要 10^4 GPU hours 來執行。以 DARTs 為代表的 Differentiable NAS 雖減少了計算需求,可隨著搜尋空間線性增長的記憶體成了新的瓶頸。由於這些侷限,之前的 NAS 都利用了 Proxy 任務,例如僅訓練少量 Epoch,只學習幾個 Block,在較小的資料集上搜尋(e.g. CIFAR)再遷移。這些 Proxy 任務上最佳化的結構,並不一定在目標任務上是最佳的。同時為了實現可轉移性,這些方法往往僅搜尋少數結構 Block,然後重複堆疊。這限制了塊的多樣性,並導致效能上的損失。依賴 Proxy 同時也意味著無法在搜尋過程中直接權衡延遲等硬體指標。

在這項工作中,我們提出了一個簡單而有效的方案來解決上述限制,稱為 ProxylessNAS:它直接在目標任務和硬體上學習結構而不依賴於 Proxy。我們還移除了先前 NAS 工作中的重複塊的限制:所有 stage 都可以自由的選擇最適合的模組。此外,為了直接在目標硬體上學習專用網路結構,在搜尋時我們也考慮了硬體指標(例如延遲)。

  • ProxylessNAS 是第一個在沒有任何代理的情況下直接學習大規模資料集(例如 ImageNet)上的 CNN 結構的 NAS 演算法,同時仍然允許大的候選集並消除重複塊的限制。它有效地擴大了搜尋空間,實現了更好的效能。

  • 我們為 NAS 提供了一種新的路徑級剪裁視角,顯示了 NAS 與模型壓縮之間的緊密聯絡(Han et al。,2016)。我們透過使用路徑級二值化將記憶體消耗節省一個數量級。

  • 我們提出了一種新的基於梯度的方法(作為一個正則函式),來處理硬體目標(例如延遲)。針對不同的硬體平臺:CPU / GPU / FPGA / TPU / NPU,ProxylessNAS 實現了針對目標硬體 CNN 結構定製。據我們所知,這是第一個來研究不同硬體結構下的專用神經網路結構的文章。

  • 廣泛的實驗證明了 Directness 和 Specialization 的優勢。它在不同硬體平臺(GPU,CPU 和 Mobile)延遲限制下,在 CIFAR-10 和 ImageNet 上實現了最好的效能。我們還分析了專用於不同硬體平臺的高效 CNN 模型的偏好,指出不同硬體平臺需要不同的神經網路結構。

方法

NAS 的路徑級減枝視角

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

我們將 NAS 建模成一個路徑級減枝的過程。如圖 2 所示,我們首先構建一個過度引數化網路。這個網路在各個位置不是採用一個確定的操作,而是保留了所有可能的操作。我們引入架構引數來顯示地學習各個位置上哪些操作是冗餘的,然後透過去除這些冗餘的操作,我們就得到了一個輕量的網路結構。如此,我們就只需要在網路結構搜尋的過程中,訓練這個過度引數化網路即可,而不需要訓練成千上萬個網路,也不需要使用一個額外的元控制器。

然而,直接去訓練這樣一個過度引數化網路是有問題的,因為其 GPU 視訊記憶體會隨著候選操作的數量線性增長。這裡,我們利用到路徑級二值化的思想來解決這個問題: 即將路徑上的架構引數二值化,並使得在訓練過程中只有一個路徑處於啟用狀態。這樣一來 GPU 視訊記憶體的需求就降到和正常訓練一個水平。在訓練這些二值化的架構引數的時候,我們採用類似 BinaryConnect 的思想,使用對應的 Binary Gate 的梯度來更新架構引數:

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

最佳化不可導的網路結構硬體指標

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

除了準確率之外,在設計高效神經網路結構的時侯,延遲(Latency)是另一個非常重要的目標。與可以使用損失函式的梯度最佳化的準確率不同,延遲這一指標是不可微的。在本節中,我們提出了兩種演算法來處理這種不可微分的目標。

如上圖所示,我們將網路結構的延遲建模為關於神經網路的連續函式。透過引入一個新的延遲損失,我們可以直接使用梯度來對其進行最佳化。

此外,作為 BinaryConnect 的替代方案,我們也可以利用 REINFORCE 來訓練這些二值化架構引數,這樣也能很自然得處理這些不可微分的目標。

結果

我們在 CIFAR-10 和 ImageNet 上進行了實驗。不同於之前的 NAS 工作,我們直接在目標資料集上進行神經網路結構學習,為目標硬體(CPU / GPU / Mobile)進行最佳化,同時允許每一個 block 自由地選擇操作。

CIFAR-10 上與之前 SOTA(state-of-the-art)模型的對比總結在表 1 中。和他們相比,我們的模型不僅在測試錯誤率上更低,並且需要的引數量也更少。例如為了達到 2.1% 的錯誤率,AmoebaNet-B 使用 34.9M 引數,而我們的模型僅使用六分之一 (5.7M 引數)。這些結果證明了直接探索大型搜尋空間而不是重複相同塊的好處。

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

ImageNet 上,我們主要關注於學習高效的 CNN 結構。因此,這是一個多目標的 NAS 任務,我們需要同時兼顧效能與延遲(不可導的硬體指標)。我們在三種不同的硬體平臺上(CPU / GPU / Mobile)進行了實驗。實驗所用的 CPU 是 2 x 2.40GHz Intel(R)Xeon(R)CPU E5-2640 v4 批大小 1, GPU 是 NVIDIA® Tesla® V100 批大小 8,Mobile 是 Google Pixel 1 單核無量化批大小 1。

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

在移動端,與 MobilenetV2 相比,我們的 ProxylessNAS 在維持同等的延遲的前提下,TOP-1 準確率提升了 2.6%。此外,在各個不同的延遲設定下,我們的模型始終大幅優於 MobilenetV2:為了達到 74.6% 的精度,MobilenetV2 需要 143ms 的推理時間,而我們模型僅需要 78ms(1.83x 倍)。與 MnasNet 相比,我們模型在提升 0.6% Top-1 的同時保持略低的推理時間。值得一提的是,我們所消耗的搜尋資源要比 Mnas 少得多 (1/200)。麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

除了移動端,我們還應用 ProxylessNAS 為 CPU 和 GPU 定製 CNN 模型。和之間 SOTA 網路結構相比,ProxylessNAS 依舊錶現出了更優的效能:在延遲低 20% 的情況下,ImageNet 上 top-1 準確率相比 MobileNetV2 提升了 3.1%。

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

之前人們習慣於將一種同一種網路結構應用到多個平臺上,而我們的實驗結果表明,我們實際上需要為不同的平臺定製神經網路結構:針對 GPU 最佳化的模型在 CPU 和移動端上執行速度並不快,反之亦然。

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

下圖展示了我們在三個硬體平臺上搜尋到的 CNN 模型的詳細結構:GPU / CPU / Mobile。我們注意到,當針對不同平臺時,網路結構呈現出不同的偏好:(i)GPU 模型短而寬,尤其是在 feature map 較大時;(ii)GPU 模型更喜歡大 MBConv 操作(例如 7x7 MBConv6),而 CPU 模型則傾向於小操作。這是因為 GPU 比 CPU 有更高的並行度,因此它可以更好地利用大 MBConv。另一個有趣的觀察是,當特徵地圖被下采樣時,所有的網路結構都傾向於選擇一個更大的 MBConv。我們認為這可能是因為大 MBConv 操作有利於網路在下采樣時保留更多資訊。值得注意的是,這是之前強制 block 之間共享結構的 NAS 方法無法發現的。

麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構

PS:我們視覺化了,網路結構隨著搜尋而變化的趨勢,影片在下鏈中。

https://hanlab.mit.edu/files/proxylessNAS/visualization.mp4

NIPS 大會期間,韓松教授將於在 workshop 分享更多本工作的資訊,有興趣的同學 / 老師歡迎前來交流。

  • 12.7 9:30AM @ Room 512 Compact Deep Neural Networks with industrial applications

  • 12.7 4:00PM @ Room 514 Machine Learning on the Phone and other Consumer Devices.

  • 12.8 9:35AM @ Room 510 Machine Learning for Systems workshop

相關文章