[ICLR2020] NAS論文導讀：Once for All: Train One Network and Specialize it for Efficient Deployment

[論文連結]

Once for All: Train One Network and Specialize it for Efficient Deployment

作者：Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han

1. Introduction

現今的DNN在各種機器學習領域都已達到了最好的成績。然而在追求準確率的背後導致的模型尺寸不斷成長，運算複雜度不斷提高的問題，也對於在不同的硬體上有效率的執行產生了很大的挑戰。例如，一個手機的程式需要可以支援多種款式的手機（目前最新的iphone或是好幾年前的iphone6），適應各種不同的硬體限制（記憶體，運算單元，cpu），甚至在同一個硬體上，也需要可以支援不同的電量以及不同的工作負擔。

但是，要設計一個可以針對不同場景所使用的網路是很花費時間和人力的，即使是現在的AutoML方法，也需要不斷的重新搜索不斷的重新訓練，因此在這樣的情況下，搜索的時間花費會隨著場景的數量線性成長，使得過去的NAS並沒有辦法應付如此大量且多元的硬體環境。

因此在這個論文中，提出了once-for-all network，可以在不同的硬體限制下，直接部屬，不用重新訓練不用重新搜索，只需要查詢表格的方式便可以決定網路的架構，且可以彈性的應對不同的depths，widths，kernel size和resolutions。

然而，要訓練這個once-for-all network並不是件容易的事情，因為訓練這個網路其實等同於訓練非常多的子網路(在這篇論文中，為10¹⁹個子網路)，而利用weight sharing的方式，雖然可以使得我們不用訓練10¹⁹個網路，但是要在weight sharing的情況下讓它們之間不會互相干擾更加的困難。因此在這篇論文中，作者提出了progressive shrinking algorithm的方式來訓練once-for-all network。

Progressive shrinking algorithm

在一開始，會訓練擁有最大的depth, width和kernel size的網路，並progressively的訓練較小的網路，使得once-for-all network可以在不會損失最大架構的once-for-all網路的情況下也可以支援較小的子網路，而這個方式也使得我們可以讓子網路在訓練時有一個好的初始值，使得我們在訓練子網路的時候不需要重頭慢慢訓練。

2. Method

2.1 Problem Formalization

將訓練once-for-all network公式化成可以支援各種不同的架構，最終的目的是希望訓練once-for-all network可以同時支援所有的子網路，並使所有的子網路達到相同程度的收斂（如同獨立的個別訓練所有的子網路）

1. W_o 為Once-for-all network的參數
2. arch_i 為architectural configurations

因此C(W_o, arch_i)代表的意思是從once-for-all network W_o 分出一步的架構並且組合成arch_i

在這篇論文當中，作者探索了四種對於convolutional neural network很重要的因素，i.e., depth, width, kernel size, and resolution

2.2 Training the Once-for-all Network

Preliminary

一個CNN通常會包含數個階段（一個階段代表的是在相同解析度的blocks）

1. Network-level：允許模型可以在不同的輸入尺寸下執行(i.e., elastic resolution)
2. Stage-level：允許每個stage可以跳過不同數量的blocks(i.e., elastic depth)
3. Block-level：允許每個block可以使用不同的channels(i.e., elastic width)，以及不同的kernel sizes(i.e., elastic kernel size)

透過不同的方式組合，在這個論文當作可以允許10¹⁹種組合，作者認為也因為有這麼多的組合，才可以針對許許多多不同的情況可以找出最符合的網路架構，但也因為有如此多種的組合，因此訓練彈性的once-for-all network便成為了一個大挑戰，首先，訓練的時間不能太長，再來，不同的子網路不能互相干擾，接下來便要介紹作者所提出的progressive shrinking approach來解決上述的問題。

A Progressive Shrinking Approach

由於直接訓練once-for-all network使得once-for-all network可以支援所有的子網路是一件很難的事情，因此作者提出了將訓練一個問題拆解成訓練一連串的子網路。

在一開始在elastic resolution的情況下訓練整個once-for-all network(最大的width, depth, kernel size)，並fine-tune once-for-all network使得網路可以同時支援最大的[width, depth, kernel size]以及部份的[width, depth, kernel size]，也就是透過progressive的方式（從最大的網路到最小的子網路）。

Progressive shrinking scheme 提供了三個優點，

1. 由於每個子網路比起整個網路可容易訓練，因此once-for-all network可加容易的收斂。
2. 先訓練大的模型再訓練小的模型使得，小的模型可容易訓練，透過保留較大模型中重要的參數，可以使得子網路有很好的初始值，這樣也使得在訓練較小的子網路的時候可以不用從頭訓練。

1. Progressive shrinking透過先訓練大的模型再慢慢訓練小型的模型可以避免子網路會使得大網路的表現變差。

Elastic Resolution

理論上，我們可以提供任何一種尺寸的照片給CNN，因為照片的尺寸並不會影響CNN的參數，但是實際上，假如在訓練時，模型並沒有看過這種尺寸的照片，他的準確度會大幅度的降低，因此為了支援elastic resolution，在訓練的時候，我們會隨機將照片放大或是縮小。

Elastic Kernel Size

理想的情況下，7*7的convolution kernel的中間區塊是可以被當作是5*5的convolution kernel，而5*5的convolution kernel的中間區塊是可以被當作是3*3的convolution kernel。但這件事實際上會遇到一個問題，也就是中間被共享的參數是需要扮演多個角色的(需要扮演7*7, 5*5以及3*3的kernel)，而對於不同的角色，中間共享的參數必然會擁有不同的分佈或是特徵，而強制的使它們在扮演不同角色時使用相同的參數會使網路的準確度變差。因此在這篇論文中，作者利用了kernel transform matrices來進行kernel間的weight sharing，也就是使用不同的kernel時，會先將原先的參數乘上kernel transform matrices，將參數的分佈轉換成較符合不同kernl的分佈。

Elastic Depth

Elastic Depth支援stage-level，也就是一連串相同output resolution的blocks。舉例來說，原先在同一個stage有N個block的once-for-all network，任何子網路都可以從這N個block中取出D個block來使用，而剩下的N-D個block便會直接跳過，因此這前D個block在最後會同時被較大的子網路以及較小的子網路共用。

Elastic Width

Width代表的是channel size的數量，也就是允許每一層可以彈性的選擇不同的expansion ratio。遵循著 progressive shrinking scheme，在一開始會訓練擁有最大的width的網路。而作者使用了channel sorting operation來確認不同channel之間的重要程度，而實做的方式是計算每一層weight的L1 Norm，較大的代表這個channel是比較重要的，舉例來說，假如要從4-channel-layer切割成3-channel-layer，便會選擇最重要的3個channel-layer（也就是L1 Norm最大的前3個），透過這個方式，子網路會被初始化會最重要的channel。

Knowledge Distillation

而在訓練once-for-all network時，會同時使用hard labels和soft labels來進行訓練，loss function如下：

2.3 Specialized Model Deployment with Once-for-all Network

當我們訓練好了once-for-all network之後，接下來便是需要能夠根據不同的需求來取出最好的子網路了，而目標是可以在符合資源需求下(e.g., latency, energy)的同時，可以最佳化準確率。而once-for-all將訓練網路和搜索網路架構拆開來，因此在這個階段是完全不需要任何的訓練成本的。

正常來說，once-for-all network可以和任何的search algorithm結合，像是RL，evolutionary algorithm，gradient descent等等。但是這些的搜索演算法需要對於每個不同的場景都需要搜索時間，導致搜索時間會隨著搜索的數量線性的成長，因此作者提出一個簡單的方式來解決搜索時間線性成長的問題，也就透過隨機從once-for-all network取樣出許多的子網路並將它們的準確率以及其他硬體的資源建各構出一個表格，因此在給定指定的資源需求的時候，便可以直接從表格中得到最符合需求的子網路，由於搜索表格的時間非常的小，因此可以節省掉隨著搜索成長的時間成本。

而作者從once-for-all network中sample了16K個子網路建構準確度表格，而由於once-for-all network需要可以支援elastic resolution，因此相同的子網路會同時測量多個不同尺寸的輸入。而作者們發現子網路的準確度會隨著輸入尺寸的變大而平緩的成長，因此作者在建構表格時，並沒有將所有的輸入尺寸都進行測量，而是以stride 16的方式進行測量(e.g., 160, 176)，而未被測量的尺寸(e.g., 164)，則會透過這個公式進行預測：

因為所有的子網路的參數都是從once-for-all切過來的，因此可以省下訓練子網路的時間，而建構表格雖然需要200 GPU hours，但是這只需要付出一次。

3. Experiments

首先，作者先在ImageNet dataset上面訓練once-for-all network，並再將once-for-all network根據許多不同的硬體設備 (Samsung Note8, Google Pixel1, Pixel2, NVIDIA 1080Ti, 2080Ti, V100 GPUs,
and Intel Xeon CPU)的資源限制切出最適合的子網路。

3.1 Training the Once-for-all Network on ImageNet

Traingin Detail

為了公平比較，作者們使用了和ProxylessNAS相同的search space。並訓練了整個once-for-all network，使之可以同時支援elastic depth，elastic width和elastic kernel size。在這樣的Search Space下，整個once-for-all network大約可以切割出2*10¹⁹個子網路。此外，once-for-all network也可以支援不同大小的輸入，以stride 4的方式從128到224。

Results

這邊列出的表格是作者證明透過Progressive shrink的方式，子網路的準確度不會比獨立訓練子網路的準確度還低，甚至會比較高。

3.2 Specialized Sub-networks for Different Hardware Platforms and Constraints

作者將訓練好的once-for-all network根據不同的硬體設備和資源限制，得到了最好的網路架構。

Comparison with NAS on Mobile

跟過去其他NAS的方法進行比較，Once-for-All更加的有效率，因為對於不同的情況下，Once-for-All都只需要constant的訓練時間和搜索時間，而有利用Progress shrink algorithm使得Once-for-All在沒有重新訓練的情況下可以跟其他的方法得到差不多的準確率，甚至更高，而在進行了25個epochs的fin-tune之後，準確率甚至可以達到75.3。

Results under Different Efficiency Constraints

而在mobile devices的表現上，在相同的速度要求下，Once-for-All也可以達到比其他NAS更高的準確率。

4. Conclusion

這篇論文提出了Once for All(OFA)，一個新的NAS方法，將訓練模型和搜索模型拆開來，使得在非常大量的架構需求下也可以很有效率的針對不同情況找到最好的網路模型。不同於以往對於每個需求都需要重新搜索重新訓練的NAS，once-for-all network可以同時支援非常多樣的架構設定，並減少非常大量的GPU訓練成本。而為了不讓once-for-all network互相干擾，提出了progressive shrinking algorithm，使得每個子網路可以達到和獨立訓練差不多的準確度。