From a87f966efd719a1e6ffbc3e2862f7c30d31d1ffa Mon Sep 17 00:00:00 2001 From: wondergo2017 Date: Wed, 8 Jun 2022 20:50:55 +0800 Subject: [PATCH] add agnnrl --- autogl/module/nas/algorithm/agnn_rl.py | 200 +++++++++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 autogl/module/nas/algorithm/agnn_rl.py diff --git a/autogl/module/nas/algorithm/agnn_rl.py b/autogl/module/nas/algorithm/agnn_rl.py new file mode 100644 index 0000000..b34c50a --- /dev/null +++ b/autogl/module/nas/algorithm/agnn_rl.py @@ -0,0 +1,200 @@ +from .rl import * +import numpy as np +class AGNNReinforceController(ReinforceController): + def resample(self,search_fields,selection): + # search_fields act as group of fields in the paper (like activation group) + self._initialize() + result = selection.copy() + + # 1. update initial state with fields not searched + for field in self.fields: + if field not in search_fields: + self._update_state(field,selection[field.name]) + + # 2. get probability of field to search + for field in search_fields: + result[field.name] = self._sample_single(field) + return result + + def _update_state(self,field,sampled): + # use unsearched fields as rnn history to update state + self._lstm_next_step() + self._inputs = self.embedding[field.name](torch.LongTensor([sampled]).to(self._inputs.device)) + +class AGNNActionGuider(nn.Module): + def __init__(self, fields, groups, **controllargs): + super(AGNNActionGuider, self).__init__() + # create independent controllers for each group + controllers=[AGNNReinforceController(fields,**controllargs) for group in groups] + self.controllers=nn.ModuleList(controllers) + self.fields=fields + self.groups=groups + + def dummy_selection(self): + # create dummy selection + result=dict() + for field in self.fields: + result[field.name]=0 + return result + + def resample(self,selection): + entropys=[] + new_selections=[] + sample_probs=[] + for idx,cont in enumerate(self.controllers): + cont=self.controllers[idx] + group=self.groups[idx] + new_selection=cont.resample(group,selection) + new_selections.append(new_selection) + entropy=cont.sample_entropy + entropys.append(entropy) + sample_probs.append(cont.sample_log_prob) + # use the most uncertain one + idx=np.argmax(entropys) + group=self.groups[idx] + print(f'$$select group {group}') + new_selection=new_selections[idx] + self.sample_log_prob=sample_probs[idx] + self.sample_entropy=entropys[idx] + print(f'$$new selection {new_selection}') + return new_selection + +@register_nas_algo("agnn") +class AGNNRL(GraphNasRL): + def search(self, space: BaseSpace, dset, estimator): + self.model = space + self.dataset = dset # .to(self.device) + self.estimator = estimator + # replace choice + self.nas_modules = [] + + k2o = get_module_order(self.model) + replace_layer_choice(self.model, PathSamplingLayerChoice, self.nas_modules) + replace_input_choice(self.model, PathSamplingInputChoice, self.nas_modules) + self.nas_modules = sort_replaced_module(k2o, self.nas_modules) + + # to device + self.model = self.model.to(self.device) + # fields + self.nas_fields = [ + ReinforceField( + name, + len(module), + isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1, + ) + for name, module in self.nas_modules + ] + + # create groups + tags='op in act concat'.split() + groups={tag:[] for tag in tags} + for field in self.nas_fields: + for tag in tags: + if tag in field.name: + groups[tag].append(field) + groups=list(groups.values()) + + # controller + self.controller = AGNNActionGuider( + self.nas_fields, + groups, + lstm_size=100, + temperature=5.0, + tanh_constant=2.5, + **(self.ctrl_kwargs or {}), + ) + self.ctrl_optim = torch.optim.Adam( + self.controller.parameters(), lr=self.ctrl_lr + ) + + # init selection (acc,selection) + self.best_selection=[0,self.controller.dummy_selection()] + + # train + with tqdm(range(self.num_epochs), disable=self.disable_progress) as bar: + for i in bar: + l2 = self._train_controller(i) + bar.set_postfix(reward_controller=l2) + + selection=self.export() + + # selections = [x[1] for x in self.hist] + # candidiate_accs = [-x[0] for x in self.hist] + # # print('candidiate accuracies',candidiate_accs) + # selection = self._choose_best(selections) + arch = space.parse_model(selection, self.device) + print(selection,arch) + return arch + + def _train_controller(self, epoch): + self.model.eval() + self.controller.train() + self.ctrl_optim.zero_grad() + rewards = [] + selections=[] + # baseline = None + baseline=self.best_selection[0] + # diff: graph nas train 100 and derive 100 for every epoch(10 epochs), we just train 100(20 epochs). totol num of samples are same (2000) + with tqdm( + range(self.ctrl_steps_aggregate), disable=self.disable_progress + ) as bar: + for ctrl_step in bar: + self._resample() + selections.append(self.selection.copy()) + metric, loss, hardware_metric = self._infer(mask="val") + reward = metric + + # bar.set_postfix(acc=metric,loss=loss.item()) + LOGGER.debug(f"{self.arch}\n{self.selection}\n{metric},{loss}") + # diff: not do reward shaping as in graphnas code + if ( + self.hardware_metric_limit is None + or hardware_metric[0] < self.hardware_metric_limit + ): + self.hist.append([-metric, self.selection]) + if len(self.hist) > self.topk: + self.hist.sort(key=lambda x: x[0]) + self.hist.pop() + rewards.append(reward) + + if self.entropy_weight: + reward += ( + self.entropy_weight * self.controller.sample_entropy.item() + ) + + if not baseline: + baseline = reward + else: + baseline = baseline * self.baseline_decay + reward * ( + 1 - self.baseline_decay + ) + + loss = self.controller.sample_log_prob * (reward - baseline) + self.ctrl_optim.zero_grad() + loss.backward() + + self.ctrl_optim.step() + + bar.set_postfix(acc=metric, max_acc=max(rewards)) + + # conserative explorer: update the best selection + idx=np.argmax(rewards) + best_reward=rewards[idx] + best_selection=selections[idx] + if best_reward>self.best_selection[0]: + self.best_selection=[best_reward,best_selection] + + print(f'$$best selection: {self.best_selection}') + LOGGER.info("epoch:{}, mean rewards:{}".format(epoch, sum(rewards) / len(rewards))) + return sum(rewards) / len(rewards) + + def _resample(self): + result = self.controller.resample(self.best_selection[1]) + self.arch = self.model.parse_model(result, device=self.device) + self.selection = result + + def export(self): + # self.controller.eval() + # with torch.no_grad(): + # return self.controller.resample() + return self.best_selection[1]