省流總結
Retiarii意思是 網斗士，是古羅馬以仿漁夫裝備——手拋網（rete）、三叉戟（fuscina 或 tridens）和匕首（pugio）進行搏斗的角斗士。官方的解釋是因為神經網絡可以看成是一個網，然后Retiarii就是去搜索網的斗士，哈哈哈，有點意思~~
Retiarii的主要貢獻有三點：
- 提供了搜索空間和搜索策略的接口，我們可以基於這個接口靈活地設計我們自己的算法，而且這兩個接口功能是解耦的，也就是說你定義好一個搜索空間后，可以無痛切換不同的搜索策略，而早期很多算法搜索空間和搜索策略是強耦合在一起的，牽一發而動全身，魔改起來很讓人頭痛
- 實現了Just-In-Time（JIT） 引擎。Retiarii框架下模型會被表示成graph，JIT引擎的作用就是將graph實例化為模型，收集搜索過程中產生的信息，並執行指定的策略等
- 實現了一系列跨模型（cross-mdoel）的優化方法來提高搜索效率。因為權重共享是在實現NAS時比較常用的一種策略，也就是說搜索空間內的子模型彼此之間有不少模塊都是overlap的，對於共享的模塊我們顯然沒必要重復定義和訓練，所以Retiarii對這個需求做了優化。
1. 前言
神經網絡架構搜索(NAS, Neural Architecture Search)儼然成為目前最熱門的AI研究方向之一，它的目的是讓算法去找到最優的模型結構，將人們從繁瑣的調參和試錯中解救出來。如果你對NAS技術還不太熟悉，建議你參考閱讀一下我之前的文章AutoML綜述。
NAS領域比較經典的工作有ENAS（基於強化學習），DARTS（基於梯度下降），AmoebaNet（基於進化算法）等等，可是這些算法的實現各不相同，想要把這些算法移植到自己的任務上難度不小，而且辛辛苦苦地魔改了很多地方還不一定能跑起來，總之這就是目前NAS研究的一個窘境。根本原因是因為現有的Deep learning框架，如Pytorch和TensorFlow本身對NAS這個需求的支持不是很好。
因此，有不少工作嘗試自動化來減輕工作量，例如AutoSklearn，TPOT，H2O.ai, Auto-Weka，Google Vizier及其開源實現Advisor,這些主要是對傳統的機器學習算法（例如SVM，決策樹等）做超參數搜索。實現對神經網絡架構進行搜索的框架有 Auto-Keras，亞馬遜的AutoGluon，華為的Vega。
本篇文章主要介紹微軟發表在OSDI2020的工作 《Retiarii: A Deep Learning Exploratory-Training Framework》，現已集成到其開源框架NNI。
2. 現有框架的痛點
痛點一：搜索空間的設計
下圖總結了現有NAS算法中涉及到的不同搜索空間的設計方式，分成了三大類
- 第一類是對單個或多個節點（節點也可以是一個復雜的模塊）做替換操作，或者增減某個節點的輸入
- 第二類是比較常見的cell結構，它和上面的替換其實有點類似，只不過差別在於這種設計方式可以使得子網絡之間共享權重，但是顯然靈活性沒有第一類高
- 第三類也可以叫做network morphism的操作，無中生有，暗度陳倉。不過新的節點的生成是要以來一定規則的，相關原理可以閱讀這篇論文 Net2Net: Accelerating Learning via Knowledge Transfer
目前很多開源的NAS代碼在實現各自的搜索空間時，以pytorch框架為例，都會在__init__和forward函數中添加邏輯控制語句，最后實現的模型非常臃腫。比如基於權重共享策略的ENAS和DARTS在搜索過程中雖然可能每次只更新某個比較小的子模型，但是整個supernet都放在了計算圖中，而且由於前向傳播過程中有很多邏輯控制，導致即使很小的模型跑起來也很耗時。
這個我自己之前跑實驗也深有體會，DARTS的某個子模型的大小比resnet18小不少，但是同樣跑一個batch的數據，resnet18就快不少。如果我沒理解錯的話，Retiarii的解決辦法就是每次會實例化一個小模型（下圖中的single model），而不是通過邏輯控制語句模擬了一個小模型（下圖中的Jumbo model）。
痛點二：搜索策略和搜索空間高度耦合
如果你讀過ENAS和DARTS的源代碼你會發現，他們設計的搜索算法是與模型高度耦合的，雖然理論上他們的算法思路適用於很多模型的搜索，但是他們設計的接口不具備較強的通用性，Retiarii則嘗試解決這一問題。
3. Mutator
我們知道現有框架定義DNN其實可以等加成數據流圖，即data-flow graph (DFG)， Retiarii也遵循這一思路，即每個節點（node）表示一個操作或者是一個子圖，而每條邊（edge）是有方向的，其表示數據的流動方向。
在Retiarii中，搜索初始時刻的模型定義為base model，如下圖示。其中的maxpool節點被定義為Mutator類，其可以變化成一個Inception cell，而這個cell本身也可搜索的，因為我們知道Inception模塊有多條路徑，那么每條路徑上的結構是什么可以通過搜索來確定。
所以，搜索空間={base models + 對base models做mutation后產生的模型集合A + 對集合A做mutation后得到的模型 集合B + ...}，換句話說隨着不斷搜索，base model的數量會從最初的一個變得越來越多,所以后面加了復數。
下面代碼中給出了Mutator的一些基礎操作，包括創建和刪除節點，連接，更新節點參數，以及從多個candidates中選擇一個或多個操作。
create_node(name:str ,op:Op,params:dict={})
delete_node(node:Node)
connect(src:NodeOutput ,dst:NodeInput)
del_connect(src:NodeOutput ,dst:NodeInput)
update_node(node:Node ,op:Op=None ,params:dict={},inputs:list=None)
choose(candidates:list ,n_chosen:int=1,type:str="choice",ctx:dict=None)
下面給出了Inception Cell的代碼實現，可以看到其需要傳入兩個參數，paths_range表示邊的數量范圍，這是可搜索的，以及每條邊上連接的節點的候選操作candidates。可以看到通過mutate函數實現了搜索和替換操作。
# define the graph mutation behavior
class InceptionMutator(BaseMutator):
def __init__(self , paths_range , candidate_ops):
self.paths_range = paths_range # [2, 3, 4, 5]
self.ops = candidate_ops # {conv , dconv , ...}
def mutate(self , targets):
if not three_node_chain(targets):
return err
n = choose(candidates=self.paths_range)
delete_node(targets[1])
for i in range(n): # create n paths
op = choose(candidates=self.ops)
nd = create_node(name=’way_’+str(i), op=op)
connect(src=targets[0].output , dst=nd.input)
connect(src=nd.output , dst=targets [2].input)
# mutation applied to the graph
apply_mutator(targets=["model/relu", "model/maxpool", "model/dense"],
mutator=InceptionMutator(
[2, 3, 4, 5], [conv , dconv , pool]))
4. JIT引擎
上圖給出了Retiarii的框架的架構示意圖，主要分成三個部分
- Input： Retiarii論文中將NAS理解為探索訓練（Exploratory-training），第一部分就是輸入，這個輸入表示的是用戶自定義的模型搜索空間（Mutator，Base Model） 以及 模型搜索策略（Model Exploration Strategy）
- JIT引擎是第二部分，也是整個框架的核心。它的運行思路是這樣的
- 1）Instantiation Control：JIT首先根據用戶指定的Exploration策略來確定本輪需要做mutation的目標模型（target model），以及需要采取什么樣的
Mutator操作。上一節的代碼中apply_mutator函數表示目標模型為["model/relu", "model/maxpool", "model/dense"]，需要執行的變異操作是InceptionMutator。 - 2）Choice Suggestion：這個其實就是
Mutator類里面定義的mutate函數。通過(1)和(2)步驟，我們可以找到指定的Target Model，並對其Apply Mutators,進而生成若干個原始的DFG。 -
- Cross-Model Optimization: 這個的作用是對raw DFGs做優化，具體細節在后面一個小結介紹，你只需要知道這樣做的目的是提高效率
- 1）Instantiation Control：JIT首先根據用戶指定的Exploration策略來確定本輪需要做mutation的目標模型（target model），以及需要采取什么樣的
- Model Training：JIT中的Training Control模塊會管理生成的若干個實例化模型的訓練過程，即解決可能存在的性能瓶頸，如CPU、GPU利用率問題
5. Cross-Model Optimization
前面介紹了，Retiarii是基於base models執行mutation操作來不斷搜索不同的模型結構，那么很顯然這些模型其實有很大一部分的結構（子圖）是相同的。一般來說，找出兩個graph的最大公共sub-graph是一個NP-hard問題，但是由於Retiarii中的所有Mutator都被定義了唯一的索引值或者是key值，所以可以很輕易地找到兩個模型之間的最大子圖。Retiarii主要從以下三個方面來對公共子圖做了優化。
5.1 Common Sub-expression Elimination (CSE)
公共子表達式消除（CSE）優化算法常用來消除一個程序中的相似操作，下面舉個簡單的例子幫助理解，假如我們在代碼里寫了這么一串計算流程：
a = 1+2+3/8+6+25
b = 15+1+2+3/8+6
c = a + b
仔細觀察可以看到變量a和b之間的公共子表達式是1+2+3/8+6,所以CSE要實現的目的就是把這個公共部分抽離出來，盡可能只執行一次，即
cse = 1+2+3/8+6
a = cse + 25
b = 15 + cse
c = a + b
這樣1+2+3/8+6就只需要執行一次即可。
那么對應到NAS任務，我們知道子網絡之間除了結構會有overlap的地方以外，數據預處理等操作基本上是一樣的，所以Retiarii則通過CSE優化算法將子網絡之間所有non-trainable的操作都merge在一起，因此只執行一次，從而大大提高了效率。注意，那些trainable的部分，比如重復的網絡結構是沒辦法merge的，因為權重需要更新。
5.2 Operator Batchting
還有一個做跨模型優化的場景就是在transfer learning的時候我們搜索模型結構，可以對不同模型的相同operator做融合。
什么意思呢？假設我們現在需要對一個新的任務場景來搜索網絡結構，但是我們已經有了一個比較好的base model，而且該模型已經有了預訓練權重。我們的目的是基於這個base model的結構和預訓練權重來搜索出新的模型結構。
假設我們通過執行mutation操作得到了如下圖示的兩個graph，其中灰色的模塊表示之前base model中已有的結構，因為是transfer learning，其權重是固定的，所以屬於non-trainable操作。
graph0和graph1的差別在於左邊有個adapter，這個adapter就表示新插入的操作，它是一個可搜索可訓練的模塊。右邊其實應該也有adapter，只不過沒畫出來而已，這里只是為了表示兩個graph絕大部分的模型架構都是一樣的。
上圖中間部分畫出了Operator Batchting的示意圖，因為graph0和1的那些灰色模塊是同樣的結構，Retiarii的做法是把兩個模型具有相同結構的部分合並在一起，我們看上圖最右邊，其實現方法是在合並后的operator的前后插入了batch和unbatch的操作，batch操作應該就是concatenate操作，即把原本兩個模型的輸入拼接在一起，這樣就將一個operator合並了，提高了計算效率；類似地，unbatch操作就是把前面計算得到的feature tensor從batch的維度划分成兩部分。
前面介紹的operator batching是假設他們的權重是相同的，例如第一個conv_3x3是有着相同權重的。其實即使權重不同也是可以做batch操作的，只需要使用比如grouped convolution或者batch_matmul等操作處理一下就可以了。
5.3 Device Placement of CSE-Optimized Graphs
另一個跨模型優化的場景則是設備放置問題。我們看看上圖就能知道是什么意思了，假設我們有4個GPU，每個GPU上對應一個Graph。這個任務預處理部分包含preprocessing(比如數據增廣)，embedding。如果embedding方法用的是word2vec，因為處理器來速度比較快，所以這個時候Retiarii會將這個操作放在CPU上執行。但是如果embedding操作是BERT網絡，那這個時候CPU計算就會成為bottleneck，所以Retiarii就會犧牲掉一個GPU專門來計算BERT embedding，然后其他GPU共用這個embedding。
這個策略和Pytorch-lightning有點類似，實現方法就是預先跑幾個iteration，profile一下iteration的時間，GPU的peak memory和利用率等信息，然后基於這些信息來分配設備和計算流程。
5.4 Super-Graph for Weight Sharing
常見的加速DNN訓練的方式主要分成兩種，一種是模型並行，另一種則是數據並行。
模型並行常見的應用場景比如說一個模型特別大，達到一個GPU都裝不下，這個時候就人為地把模型划分成若干個部分，然后每個GPU負責不同部分的計算，只不過后一個GPU每次都需要等前一個計算的結果。當然也可以通過一些schedule的方式隱藏掉等待時間，這不是本文的重點。
另一種是更常見的加速方式，即數據並行。在Pytorch中我們常用的DataParallel就屬於這種，即我們把模型copy多分到不同GPU上，然后不同GPU用不同的batch的數據進行訓練。
Retiarii的做法則是將二者融合，這一點天然地適合NAS任務，不過Retiarii論文中並沒有系統地去評估這一方法。
6. 框架評估
Retiarii設計了一系列實驗來評估它的計算性能，用到的指標是throughput (samples/s)，即每秒能處理多少個樣本數據。
6.1 Shared Data Loading and Preprocessing
前面已經提到過，NAS任務中，不同子模型其實是有着相同的數據讀取和預處理操作的。 而為了提高效率，Retiarii使用了common sub-expression elimination (CSE) 優化算法。作為對比，baseline是每個模型就像普通的DNN一樣單獨訓練，不過使用了NVIDIA提供的CUDA Multi-Process Service (MPS)。MPS的作用簡單理解就是當我們在一個GPU上執行多個計算任務，MPS能夠起到加速作用，具體地用法和原理參考官方技術文章。實驗是在4個V100上跑的，CPU的核數是20。
模型根據round-robin（循環制）的方式分發，即假如有6個模型，只有4個GPU，那個前兩個GPU上分發了2個模型，后兩個GPU只有一個模型。下圖給出了實驗結果，可以看到隨着模型數量的增加，使用的CSE的Retiarii框架的throughput基本呈現線性增長，而baseline無法scale。不過可以看到當模型數量從4擴展到5的時候，throughput沒有提升，還反而下降了，文中的解釋是第一個GPU由於需要訓練2個模型，而且Retiarii將不同模型做了merge操作，所以其他模型必須等待第一個GPU。總之下圖的結果表明CSE優化操作的確可以有效提高CPU的計算效率。
下圖展示了TextNAS的測試結果，比較的方式是將BERT Embedding放置到CPU和GPU的throughput。因為Retiarii會犧牲一個GPU用來計算BERT embedding，所以下圖中Retiarii的模型數量都是3的倍數。注意Baseline也將BERT embedding放置到了GPU上計算。
不過我們可以看到，即使少了一個GPU，但是Retiarii的計算性能還是遠遠高於baseline。
6.2 Operator Batching
說實話下面這個結果震驚到我了，這個實驗采用的是預訓練好的MobileNet，設備是單個16GB的V100。
Baseline是開啟CUDA MPS后計算多個模型；Retiarii和baseline的差別在於它對這些子模型執行了Operator Batching操作。二者采樣模型的順序是一樣的。可以看到Baseline在采樣到12個模型后就OOM了，而Retiarii能擴展到192個模型，NB！而且同樣擴展到12個模型時，Retiarii的Throughput是baseline的1.76倍。
6.3 Optimization for Weight Sharing
這個實驗使用的是SPOS算法，模型搜索空間基於ShuffleNetV2 block，每個batch會采樣生成一個新的模型實例。兩個baseline分別是：
- baseline1：supernet模型的checkpoint存到文件里，即每次從這個文件里load子模型的權重，訓練一個batch后，又把權重保存到文件里
- baseline2：Supernet的權重保存到GPU里，以
dict對象的形式儲存 - Retiarii的做法則是會收集所有已采樣的模型結構，並且會把這些結構merge成一個supernet，注意這和傳統的方式不一樣，傳統的方法是預先手動定義好Supernet的結構，而Retiarii是在搜索過程中不斷更新Supernet的結構和權重。
6.4 搜索速度
這個實驗設置了兩個baseline來比較搜索速度，簡單來說就是采樣1000個模型花了多少時間。
實驗設置如下：
- 選擇了Mnasnet，NASNet和AmoebaNet三個NAS算法，因為前面兩個都是采樣了強化學習的搜索策略，后面兩個的搜索空間一直，俗稱變量控制法，非常合理。
- baseline 1： 還是像前面的實驗一樣，每個被采樣得到的模型都是像常規的DNN模型一樣在不同的GPU上獨立訓練
- baseline 2： 開啟CUDA MPS，如果一個GPU上有多個模型時可以起到加速作用，只不過他們的graph不會merge
- Retiarii：就是在baseline2的基礎上把模型的graph merge起來
實驗結果如下面兩個圖所示，第一個使用了NVIDIA的大力出奇跡的DALI數據讀取包，第二個是使用Pytorch自帶的數據讀取包，可以看到Retiarii的加速效果還是非常明顯的
6.5 Scaling Weight-Shared Training
前面介紹過mixed-parallelism策略，這個實驗就是為了檢測這個策略的有效性。實驗設置如下：
- 實驗使用SPOS算法基於ShuffleNet搜索，每個batch會搜索一個新的模型。
- 數據集使用的是ImageNet，總共跑了60個epoch，圖片數量是1,281,167，batch size是256。所以總共采樣訓練了1281167/256*60=300240個模型。
- 實驗是在兩個節點上運行，每個節點4個V100。
- 評估方法是隨機采樣196個模型，並使用256個驗證圖像來評估這些模型的性能。196個模型的平均驗證准確率視為算法的性能，越高越好。
下圖給出了不同策略下的性能比較，柱狀圖表示訓練的時間，折線表示對應的平均accuracy。前三個使用的都是基於常見的預先定義好的Supernet訓練，最后的Mixed Parallel則是Retiarii框架的方法。
可以看到通過混合並行的策略不僅加快了訓練時間，而且也保證了和SyncBN Data Parallel差不多的性能。
雖然Data Parallel增加batch size到2048的時候將時間減少地比Mixed模式還小，但是准確率卻大幅下降，這是因為大batch size會損害模型性能。
SyncBN 是指在多個GPU之間同步計算batch normalization的操作，這個可以增加模型質量，只不過以訓練時長增加作為代價。
