ECCV20 BigNAS無需後處理直接部署

【GiantPandaCV導語】這篇是Slimmable Network三部曲之後的續作,提出了Single-Stage的超網訓練方法,在更大的搜索空間中,提出了很多訓練的Trick來提升訓練效果以及穩定訓練過程。

0. Info

Title: BigNAS: Scaling Up Neural Architecture Search with Big Single-Stage Models

Author: Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, Quoc Le Google Brain

Link: https://arxiv.org/pdf/2003.11142v3.pdf

Date: ECCV20

Code: https://github.com/JiahuiYu/slimmable_networks

1. Motivation

目前NAS的一個非常受歡迎的研究分支是one-shot NAS,即用一個超網包含所有的搜索空間,子網絡通過共享超網絡的參數來加速訓練流程。超網的作用可以看作一個評估器,用於通過評估子網在驗證集上的準確率來代表該網絡在測試集的效果。

One-Shot NAS在這裏被劃分爲Two-Stage NAS,包含了搜索階段和評估階段。

  • 在搜索階段,通過使用某些算法進行訓練超網絡,找到最有希望的子網;
  • 在評估階段,將搜索階段找到的子網從頭進行訓練,得到最終在測試集上的準確率。

之所以需要評估階段是因爲一個實驗現象 ,即通過超網共享得到的子網的準確率是遠低於從頭訓練該子網的準確率的。所以之前的大部分工作都存在評估階段,通常會在搜索階段結束以後採用重新訓練、微調或者其他後處理髮方法提升準確率。

BigNAS則試圖挑戰以上觀點,認爲評估階段可以不必要存在的,無需額外的重新訓練或者後處理步驟,從超網中直接得到的子網絡就可以直接部署,也即提出了One-Stage NAS。

如上圖所示:普通的One-Shot模型使用訓練好的超網得到對應子網,然後經過retrain纔可以進行部署。

Once for all中則提出了Progressive Shrinking的方法,先訓練大的網絡,然後大的網絡作爲教師網絡指導小網絡進行訓練,最終可以可以無需retrain直接部署。

BigNAS則提出使用了Single-Stage的範式,同時訓練所有的子網絡,並且可以無需retrain直接部署。

2. Contribution

Once for all並沒有完全拋棄後處理過程,根據OFA提供的試驗結果,其在ImageNet上繼續finetune了25個epoch,也會帶來一定的提升。

BigNAS的核心貢獻就是:

  • 無需重新訓練和後處理,可以訓練一個Single-Stage模型,直接從超網中就可以切分子網用於部署。

  • 提出了一些技巧來彌補不同的初始化和大小模型的學習率調整的差距。

筆者個人感覺,超網絡作爲一個過參數的網絡,相當於集成了成千上萬個子網絡進行共享,每個子網絡的梯度優化方向不同,可能會帶來衝突,這可能是weight sharing nas效果不如stand alone nas的原因。那爲何BigNAS可以做到無需retrain直接部署呢?筆者覺得這與超網絡的模型容量有一定的關係,同時也和訓練優化的技巧有一定的關係。OFA得到的模型的FLOPS在230M的量級,BigNAS的FLOPS則是從200M到1G都有,可以說模型容量比較大,如果能夠承載非常多的子網,那麼可能準確率本身確實可以達到比較高的水平。再配合上一些探索得到的訓練技巧,將超網性能提升到一個非常高的水平。廖哥認爲inplace distillation起到了比較重要的作用。鑫哥認爲超網模型的參數冗餘度非常高,所以OFA這種類型是可行的。也歡迎添加筆者微信進行討論~

3. Method

方法部分主要介紹了幾種訓練技巧以及選擇模型的方式。

筆者之間訓練過一個簡單的超網,可以說超網訓練難度要比訓練普通的網絡高很多倍,由於其每次是動態採樣,可能會出現訓練不穩定,大模型過擬合小模型欠擬合的問題。

3.1 Tricks

Trick1: 三明治法則

每一步訓練,找到最大模型和最小模型,然後無偏均勻採樣N(N=2)箇中間大小的模型。訓練以上所有模型,累積梯度,然後進行一次更新。

最小的模型代表了整體的最低水平,最大的模型代表整體最高水平,BigNAS目標就是同時推高lower bound和upper bound,實現整個超網性能的提升。

Trick2: Inplace Distillation

使用最大模型提供的soft label對其他小網絡進行蒸餾也是一種常用的trick(OFA中就用到了這種策略)。

值得一提的是,BigNAS中的小網絡只使用inplace distillation進行訓練,並沒有直接根據label進行訓練。

Trick3: Initialization

訓練超網時候的learning rate選取不能和普通網絡一致,需要將learning rate調整爲原來的30%左右,這樣訓練loss纔不會爆炸。

這樣做雖然loss不會爆炸,但是也降低了最終的性能,這就需要另外一個方法來穩定訓練。由於本文搜索空間中是包含殘差網絡的,使用zero-initialize每個殘差塊的最後一個BN的縮放係數爲0,不僅可以穩定訓練還能提升最終模型準確率。

  # Zero-initialize the last BN in each residual branch,
  # so that the residual branch starts with zeros, and each residual block behaves like an identity.
  # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
  for m in self.modules():
      if isinstance(m, MaskedBlock):
          # type: ignore[arg-type]
          nn.init.constant_(m.shortcut.convbn.bn.weight, 0)

Trick4: Convergence Behavior

使用三明治法則進行訓練會帶來一個問題,隨着訓練的進行,在小網絡還處於欠擬合的時候,大網絡就已經過擬合了。

BigNAS提出修改學習率策略來解決這個問題,使用指數衰減策略降低學習率的時候,當學習率達到初始設定的5%, 將學習率保持恆定。

這種做法的理由是,大模型達到精度峯值以後用恆定學習率可以梯度震盪,緩解過擬合線性,而這個時候欠擬合的小模型可以繼續進行訓練。

Trick5: Regularization

正則化方法,也是爲了解決以上問題,只給最大的子網施加正則,包括weight decay、Dropout等正則化方法。

Trick6: BN Calibration

BN矯正技術在one-shot nas中非常常見,single path one shot、FairNAS等工作中都使用到了這個技術。採樣得到的子網通常需要在訓練集上重新計算BN的值,這樣模型性能纔不會大幅降低。

3.2 Coarse-to-fine Architecture Selection

由於現在搜索空間中的候選網絡數量非常多,如何從這些候選網絡中選擇出合適的網絡呢?BigNAS提出了由粗到細的搜索方法:

  • 粗略搜索階段:找到一個大致滿足約束的粗粒度空間,

  • 細粒度搜索階段:然後再這個子空間中進行細粒度網格搜索查找模型。

4. Experiment

  1. 搜索空間展示:使用的是MobileNetV2的搜索空間。

  1. 在不同量級下與其他方法進行對比:

  1. 驗證學習率設置以及初始化方法對模型訓練的影響。

  1. 驗證學習率策略的有效性

  1. 驗證正則方法的有效性

  1. 驗證繼續finetune是否會帶來提升?

  1. 驗證從頭訓練會不會帶來提升?

5. Revisiting

看這篇文章很容易把BigNAS看作是Trick的合集,不得不承認,BigNAS的效果確實很好,少不了這些Trick的功勞。後期筆者會對這些trick進行驗證,看看是否真的和論文中一樣work。

其中比較令人在意的一個問題就是爲何這樣可以work,即直接從超網中採樣得到子網就能直接部署,無需finetune,retrain等操作。

可能的解釋有:

  • 模型容量足夠大,之前的two stage nas算法容量不足,參數冗餘度不夠。

  • 之前的two-stage nas沒有充分訓練

  • 訓練技巧的功勞,比如Progressive Shrinking方法或者Inplace Distillation的功勞。

此外,OFA在MobileNetV3, ProxylessNAS, ResNet三種空間進行搜索;BigNAS在MobileNet,ResNet,MNasNet進行的搜索;這幾種空間其實差不太多,都是ResNet-like的類型,其他two-stage nas的空間,比如DARTS-like,NASNet-like搜索空間中是否可行呢?

6. Take Away

BigNAS追隨了Once for all的工作,提出了One-Stage NAS,無需retrain就直接部署的方案,使用一系列技巧來提升模型訓練的效果,得到了非常廣泛的BigNAS族羣,FLOPS從200M到1G都廣泛存在。

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