ENAS的原理和代碼解析

05-03

高效神經架構搜索（Efficient Neural Architecture Search，[pdf] ）是Google提出一種實用的自動化模型設計方法，它克服了神經架構搜索（Neural Architecture Search，[pdf] ）算力成本巨大的缺陷，將算力成本減少1000倍以上，僅用一塊英偉達GTX 1080Ti顯卡，就能在16個小時之內完成架構搜索。

其實架構搜索在算力成本上的早已有了突破，SMASH通過隨機採樣架構，在訓練集上通過SGD學習HyperNet權重。並使用HyperNet生成的權重初始化子網路，比每次隨機初始化權重並從頭訓練子網路減少大量算力成本。

ENAS的原理

神經網路架構的搜索空間可以表示成有向無環圖（DAG），一個神經網路架構可以表示成DAG的一個子圖。

圖片來自[1]

上圖是節點數為5的搜索空間，紅色箭頭連接的子圖表示一個神經網路架構。圖中節點表示計算，邊表示信息流。ENAS使用一個RNN（稱為controller）決定每個節點的計算類型和選擇激活哪些邊。ENAS中使用節點數為12的搜索空間，計算類型為tanh，relu，identity，sigmoid四種激活函數，所以搜索空間有 $4^N imes N！$ 約 $10^{15}$ 種神經網路架構。

controller工作流程：

以節點數為4的搜索空間為例。

controller選擇節點1的計算類型為tanh（節點1的前置節點是輸入）；選擇節點2的前置節點為1，計算類型為ReLU；選擇節點3的前置節點為2，計算類型為ReLU；選擇節點4的前置節點為1，計算類型為tanh。

便得到如下的RNN神經網路架構：節點3和節點4是葉子節點，他們輸出的平均值作為RNN神經網路架構的輸出。該神經網路架構的參數由 $w^{1,2}$ ， $w^{2,4}$ ， $w^{2,3}$ ， $w^{3,4}$ 組成，其中 $w^{i,j}$ 是節點 $i$ 和節點 $j$ 之間的參數。

在NAS中， $w^{1,2}$ ， $w^{2,4}$ ， $w^{2,3}$ ， $w^{3,4}$ 都是隨機初始化，並在每個神經網路架構中從頭開始訓練的。在ENAS，這些參數是所有神經網路架構共享的。如果下一次controller得到的神經網路架構如下，它參數由 $w^{1,2}$ ， $w^{2,4}$ ， $w^{2,3}$ ， $w^{3,4}$ 組成，其中 $w^{1,2}$ ， $w^{2,3}$ ， $w^{3,4}$ 與上面神經網路架構相同的。

通過參數共享，ENAS解決了NAS算力成本巨大的缺陷。

ENAS工作流程：

loop

controller固定參數 $heta$ ，採樣一個神經網路架構，在訓練集中訓練該神經網路架構，並通過 SGD調整神經網路架構的參數 $w$ 。

end loop

loop

controller採樣一組神經網路架構，在驗證集上計算ppl，並根據ppl和controller的交叉熵計算reward，通過Adam調整controller的參數 $heta$ 。

end loop

ENAS源碼解析

ENAS非官方的開源代碼carpedm20/ENAS-pytorch

以 Penn Treebank為例，ENAS總共訓練150epoch，trainer.py:89

for self.epoch in range(self.start_epoch, self.args.max_epoch): # 1. Training the shared parameters ω of the child models self.train_shared() # 2. Training the controller parameters θ self.train_controller()

訓練共享參數 $w$ 。訓練400個step，每個step採樣一個的神經網路架構，並在一個batch的訓練集數據上運行SGD，更新神經網路架構的參數。SGD的學習率為20，15個epoch後每個epoch以0.96衰減。trainer.py:141

while train_idx < self.train_data.size(0) - 1 - 1: if step > max_step: break dags = self.controller.sample(self.args.shared_num_sample) inputs, targets = self.get_batch(self.train_data, train_idx, self.max_length) loss = self.get_loss(inputs, targets, hidden, dags) #loss, hidden = self.get_loss(inputs, targets, hidden, dags, with_hidden=True) #hidden = detach(hidden) # update self.shared_optim.zero_grad() loss.backward() t.nn.utils.clip_grad_norm( model.parameters(), self.args.shared_grad_clip) self.shared_optim.step() total_loss += loss.data

訓練controller參數 $heta$ 。訓練2000個step。每個step採樣一組神經網路架構，並在計算reward，根據reward運行Adam，更新controller參數。Adam的學習率為0.00035。trainer.py:223

for step in trange(self.args.controller_max_step, desc="train_controller"): # sample models dags, log_probs, entropies = self.controller.sample(with_details=True) # calculate reward np_entropies = entropies.data.cpu().numpy() rewards = self.get_reward(dags, np_entropies, valid_idx) # discount if 1 > self.args.discount > 0: rewards = discount(rewards, self.args.discount) reward_history.extend(rewards) entropy_history.extend(np_entropies) # moving average baseline if baseline is None: baseline = rewards else: decay = self.args.ema_baseline_decay baseline = decay * baseline + (1 - decay) * rewards adv = rewards - baseline adv_history.extend(adv) # policy loss loss = - log_probs * get_variable(adv, self.cuda, requires_grad=False) if self.args.entropy_mode == regularizer: loss -= self.args.entropy_coeff * entropies loss = loss.sum() # or loss.mean() pbar.set_description( f"train_controller| R: {rewards.mean():8.6f} | R-b: {adv.mean():8.6f} " f"| loss: {loss.cpu().data[0]:8.6f}") # update self.controller_optim.zero_grad() loss.backward() if self.args.controller_grad_clip > 0: t.nn.utils.clip_grad_norm( model.parameters(), self.args.controller_grad_clip) self.controller_optim.step() total_loss += loss.data[0]

在驗證集數據上計算ppl，根據ppl和controller的交叉墒計算reward。計算公式為c / ppl + entropy_coeff * entropies，其中entropy_coeff是0.1。trainer.py:183

def get_reward(self, dag, entropies, valid_idx=None): if type(entropies) is not np.ndarray: entropies = entropies.data.cpu().numpy() if valid_idx: valid_idx = 0 inputs, targets = self.get_batch(self.valid_data, valid_idx, self.max_length) valid_loss = self.get_loss(inputs, targets, None, dag) valid_ppl = math.exp(valid_loss.data[0]) # TODO: we dont know reward_c if self.args.ppl_square: # TODO: but we do know reward_c=80 in the previous paper R = self.args.reward_c / valid_ppl ** 2 else: R = self.args.reward_c / valid_ppl if self.args.entropy_mode == reward: rewards = R + self.args.entropy_coeff * entropies elif self.args.entropy_mode == regularizer: rewards = R * np.ones_like(entropies) else: raise NotImplemented(f"Unkown entropy mode: {self.args.entropy_mode}") return rewards

訓練結束後，採樣100個神經網路架構，選擇reward最大的那個，就是目標解。trainer.py:338

def derive(self, sample_num=None, valid_idx=0): if sample_num is None: sample_num = self.args.derive_num_sample dags, log_probs, entropies = self.controller.sample(sample_num, with_details=True) max_R, best_dag = 0, None pbar = tqdm(dags, desc="derive") for dag in pbar: R = self.get_reward(dag, entropies, valid_idx) if R.max() > max_R: max_R = R.max() best_dag = dag pbar.set_description(f"derive| max_R: {max_R:8.6f}") fname = f"{self.epoch:03d}-{self.controller_step:06d}-{max_R:6.4f}-best.png" path = os.path.join(self.args.model_dir, "networks", fname) draw_network(best_dag, path) self.tb.image_summary("derive/best", [path], self.epoch) return best_dag

更多關於架構搜索，可以關注markdtw/awesome-architecture-search

[1] [1802.03268] Efficient Neural Architecture Search via Parameter Sharing