秦一鉴
GitHub
5 years ago
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 599 additions and 234 deletions
Split View
Diff Options
-
+44 -226autogl/module/nas/algorithm/enas.py
-
+383 -0autogl/module/nas/algorithm/rl.py
-
+26 -2autogl/module/nas/estimator/one_shot.py
-
+29 -4autogl/module/nas/space/graph_nas.py
-
+2 -2examples/test_graph_nas.py
-
+42 -0examples/test_graph_nas_rl.py
-
+73 -0examples/test_graph_nas_space.py
+ 44
- 226
autogl/module/nas/algorithm/enas.py
View File
| @@ -10,219 +10,9 @@ from .base import BaseNAS | |||
| from ..space import BaseSpace | |||
| from ..utils import AverageMeterGroup, replace_layer_choice, replace_input_choice, get_module_order, sort_replaced_module | |||
| from nni.nas.pytorch.fixed import apply_fixed_architecture | |||
| from tqdm import tqdm | |||
| _logger = logging.getLogger(__name__) | |||
| def _get_mask(sampled, total): | |||
| multihot = [i == sampled or (isinstance(sampled, list) and i in sampled) for i in range(total)] | |||
| return torch.tensor(multihot, dtype=torch.bool) # pylint: disable=not-callable | |||
| class PathSamplingLayerChoice(nn.Module): | |||
| """ | |||
| Mixed module, in which fprop is decided by exactly one or multiple (sampled) module. | |||
| If multiple module is selected, the result will be sumed and returned. | |||
| Attributes | |||
| ---------- | |||
| sampled : int or list of int | |||
| Sampled module indices. | |||
| mask : tensor | |||
| A multi-hot bool 1D-tensor representing the sampled mask. | |||
| """ | |||
| def __init__(self, layer_choice): | |||
| super(PathSamplingLayerChoice, self).__init__() | |||
| self.op_names = [] | |||
| for name, module in layer_choice.named_children(): | |||
| self.add_module(name, module) | |||
| self.op_names.append(name) | |||
| assert self.op_names, 'There has to be at least one op to choose from.' | |||
| self.sampled = None # sampled can be either a list of indices or an index | |||
| def forward(self, *args, **kwargs): | |||
| assert self.sampled is not None, 'At least one path needs to be sampled before fprop.' | |||
| if isinstance(self.sampled, list): | |||
| return sum([getattr(self, self.op_names[i])(*args, **kwargs) for i in self.sampled]) # pylint: disable=not-an-iterable | |||
| else: | |||
| return getattr(self, self.op_names[self.sampled])(*args, **kwargs) # pylint: disable=invalid-sequence-index | |||
| def __len__(self): | |||
| return len(self.op_names) | |||
| @property | |||
| def mask(self): | |||
| return _get_mask(self.sampled, len(self)) | |||
| class PathSamplingInputChoice(nn.Module): | |||
| """ | |||
| Mixed input. Take a list of tensor as input, select some of them and return the sum. | |||
| Attributes | |||
| ---------- | |||
| sampled : int or list of int | |||
| Sampled module indices. | |||
| mask : tensor | |||
| A multi-hot bool 1D-tensor representing the sampled mask. | |||
| """ | |||
| def __init__(self, input_choice): | |||
| super(PathSamplingInputChoice, self).__init__() | |||
| self.n_candidates = input_choice.n_candidates | |||
| self.n_chosen = input_choice.n_chosen | |||
| self.sampled = None | |||
| def forward(self, input_tensors): | |||
| if isinstance(self.sampled, list): | |||
| return sum([input_tensors[t] for t in self.sampled]) # pylint: disable=not-an-iterable | |||
| else: | |||
| return input_tensors[self.sampled] | |||
| def __len__(self): | |||
| return self.n_candidates | |||
| @property | |||
| def mask(self): | |||
| return _get_mask(self.sampled, len(self)) | |||
| class StackedLSTMCell(nn.Module): | |||
| def __init__(self, layers, size, bias): | |||
| super().__init__() | |||
| self.lstm_num_layers = layers | |||
| self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias) | |||
| for _ in range(self.lstm_num_layers)]) | |||
| def forward(self, inputs, hidden): | |||
| prev_h, prev_c = hidden | |||
| next_h, next_c = [], [] | |||
| for i, m in enumerate(self.lstm_modules): | |||
| curr_h, curr_c = m(inputs, (prev_h[i], prev_c[i])) | |||
| next_c.append(curr_c) | |||
| next_h.append(curr_h) | |||
| # current implementation only supports batch size equals 1, | |||
| # but the algorithm does not necessarily have this limitation | |||
| inputs = curr_h[-1].view(1, -1) | |||
| return next_h, next_c | |||
| class ReinforceField: | |||
| """ | |||
| A field with ``name``, with ``total`` choices. ``choose_one`` is true if one and only one is meant to be | |||
| selected. Otherwise, any number of choices can be chosen. | |||
| """ | |||
| def __init__(self, name, total, choose_one): | |||
| self.name = name | |||
| self.total = total | |||
| self.choose_one = choose_one | |||
| def __repr__(self): | |||
| return f'ReinforceField(name={self.name}, total={self.total}, choose_one={self.choose_one})' | |||
| class ReinforceController(nn.Module): | |||
| """ | |||
| A controller that mutates the graph with RL. | |||
| Parameters | |||
| ---------- | |||
| fields : list of ReinforceField | |||
| List of fields to choose. | |||
| lstm_size : int | |||
| Controller LSTM hidden units. | |||
| lstm_num_layers : int | |||
| Number of layers for stacked LSTM. | |||
| tanh_constant : float | |||
| Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``. | |||
| skip_target : float | |||
| Target probability that skipconnect will appear. | |||
| temperature : float | |||
| Temperature constant that divides the logits. | |||
| entropy_reduction : str | |||
| Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced. | |||
| """ | |||
| def __init__(self, fields, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, | |||
| skip_target=0.4, temperature=None, entropy_reduction='sum'): | |||
| super(ReinforceController, self).__init__() | |||
| self.fields = fields | |||
| self.lstm_size = lstm_size | |||
| self.lstm_num_layers = lstm_num_layers | |||
| self.tanh_constant = tanh_constant | |||
| self.temperature = temperature | |||
| self.skip_target = skip_target | |||
| self.lstm = StackedLSTMCell(self.lstm_num_layers, self.lstm_size, False) | |||
| self.attn_anchor = nn.Linear(self.lstm_size, self.lstm_size, bias=False) | |||
| self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False) | |||
| self.v_attn = nn.Linear(self.lstm_size, 1, bias=False) | |||
| self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1) | |||
| self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), # pylint: disable=not-callable | |||
| requires_grad=False) | |||
| assert entropy_reduction in ['sum', 'mean'], 'Entropy reduction must be one of sum and mean.' | |||
| self.entropy_reduction = torch.sum if entropy_reduction == 'sum' else torch.mean | |||
| self.cross_entropy_loss = nn.CrossEntropyLoss(reduction='none') | |||
| self.soft = nn.ModuleDict({ | |||
| field.name: nn.Linear(self.lstm_size, field.total, bias=False) for field in fields | |||
| }) | |||
| self.embedding = nn.ModuleDict({ | |||
| field.name: nn.Embedding(field.total, self.lstm_size) for field in fields | |||
| }) | |||
| def resample(self): | |||
| self._initialize() | |||
| result = dict() | |||
| for field in self.fields: | |||
| result[field.name] = self._sample_single(field) | |||
| return result | |||
| def _initialize(self): | |||
| self._inputs = self.g_emb.data | |||
| self._c = [torch.zeros((1, self.lstm_size), | |||
| dtype=self._inputs.dtype, | |||
| device=self._inputs.device) for _ in range(self.lstm_num_layers)] | |||
| self._h = [torch.zeros((1, self.lstm_size), | |||
| dtype=self._inputs.dtype, | |||
| device=self._inputs.device) for _ in range(self.lstm_num_layers)] | |||
| self.sample_log_prob = 0 | |||
| self.sample_entropy = 0 | |||
| self.sample_skip_penalty = 0 | |||
| def _lstm_next_step(self): | |||
| self._h, self._c = self.lstm(self._inputs, (self._h, self._c)) | |||
| def _sample_single(self, field): | |||
| self._lstm_next_step() | |||
| logit = self.soft[field.name](self._h[-1]) | |||
| if self.temperature is not None: | |||
| logit /= self.temperature | |||
| if self.tanh_constant is not None: | |||
| logit = self.tanh_constant * torch.tanh(logit) | |||
| if field.choose_one: | |||
| sampled = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1) | |||
| log_prob = self.cross_entropy_loss(logit, sampled) | |||
| self._inputs = self.embedding[field.name](sampled) | |||
| else: | |||
| logit = logit.view(-1, 1) | |||
| logit = torch.cat([-logit, logit], 1) # pylint: disable=invalid-unary-operand-type | |||
| sampled = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1) | |||
| skip_prob = torch.sigmoid(logit) | |||
| kl = torch.sum(skip_prob * torch.log(skip_prob / self.skip_targets)) | |||
| self.sample_skip_penalty += kl | |||
| log_prob = self.cross_entropy_loss(logit, sampled) | |||
| sampled = sampled.nonzero().view(-1) | |||
| if sampled.sum().item(): | |||
| self._inputs = (torch.sum(self.embedding[field.name](sampled.view(-1)), 0) / (1. + torch.sum(sampled))).unsqueeze(0) | |||
| else: | |||
| self._inputs = torch.zeros(1, self.lstm_size, device=self.embedding[field.name].weight.device) | |||
| sampled = sampled.detach().numpy().tolist() | |||
| self.sample_log_prob += self.entropy_reduction(log_prob) | |||
| entropy = (log_prob * torch.exp(-log_prob)).detach() # pylint: disable=invalid-unary-operand-type | |||
| self.sample_entropy += self.entropy_reduction(entropy) | |||
| if len(sampled) == 1: | |||
| sampled = sampled[0] | |||
| return sampled | |||
| from .rl import PathSamplingLayerChoice,PathSamplingInputChoice,ReinforceField,ReinforceController | |||
| class Enas(BaseNAS): | |||
| """ | |||
| @@ -272,7 +62,7 @@ class Enas(BaseNAS): | |||
| def __init__(self, device='cuda', workers=4,log_frequency=None, | |||
| grad_clip=5., entropy_weight=0.0001, skip_weight=0.8, baseline_decay=0.999, | |||
| ctrl_lr=0.00035, ctrl_steps_aggregate=20, ctrl_kwargs=None,*args,**kwargs): | |||
| ctrl_lr=0.00035, ctrl_steps_aggregate=20, ctrl_kwargs=None,n_warmup=100,model_lr=5e-3,model_wd=5e-4,*args,**kwargs): | |||
| super().__init__(device) | |||
| self.device=device | |||
| self.num_epochs = kwargs.get("num_epochs", 5) | |||
| @@ -287,14 +77,13 @@ class Enas(BaseNAS): | |||
| self.workers = workers | |||
| self.ctrl_kwargs=ctrl_kwargs | |||
| self.ctrl_lr=ctrl_lr | |||
| self.n_warmup=n_warmup | |||
| self.model_lr = model_lr | |||
| self.model_wd = model_wd | |||
| def search(self, space: BaseSpace, dset, estimator): | |||
| self.model = space | |||
| self.dataset = dset#.to(self.device) | |||
| self.estimator = estimator | |||
| self.model_optim = torch.optim.SGD( | |||
| self.model.parameters(), lr=0.01, weight_decay=3e-4 | |||
| ) | |||
| self.estimator = estimator | |||
| # replace choice | |||
| self.nas_modules = [] | |||
| @@ -305,20 +94,44 @@ class Enas(BaseNAS): | |||
| # to device | |||
| self.model = self.model.to(self.device) | |||
| self.model_optim = torch.optim.Adam( | |||
| self.model.parameters(), lr=self.model_lr, weight_decay=self.model_wd | |||
| ) | |||
| # fields | |||
| self.nas_fields = [ReinforceField(name, len(module), | |||
| isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1) | |||
| for name, module in self.nas_modules] | |||
| self.controller = ReinforceController(self.nas_fields, **(self.ctrl_kwargs or {})) | |||
| self.ctrl_optim = torch.optim.Adam(self.controller.parameters(), lr=self.ctrl_lr) | |||
| # warm up supernet | |||
| with tqdm(range(self.n_warmup)) as bar: | |||
| for i in bar: | |||
| acc,l1=self._train_model(i) | |||
| with torch.no_grad(): | |||
| val_acc,val_loss=self._infer('val') | |||
| bar.set_postfix(loss=l1,acc=acc,val_acc=val_acc,val_loss=val_loss) | |||
| # train | |||
| for i in range(self.num_epochs): | |||
| self._train_model(i) | |||
| self._train_controller(i) | |||
| with tqdm(range(self.num_epochs)) as bar: | |||
| for i in bar: | |||
| try: | |||
| l1=self._train_model(i) | |||
| l2=self._train_controller(i) | |||
| except Exception as e: | |||
| print(e) | |||
| nm=self.nas_modules | |||
| for i in range(len(nm)): | |||
| print(nm[i][1].sampled) | |||
| import pdb | |||
| pdb.set_trace() | |||
| bar.set_postfix(loss_model=l1,reward_controller=l2) | |||
| selection=self.export() | |||
| print(selection) | |||
| return space.export(selection,self.device) | |||
| def _train_model(self, epoch): | |||
| self.model.train() | |||
| self.controller.eval() | |||
| @@ -330,15 +143,19 @@ class Enas(BaseNAS): | |||
| nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip) | |||
| self.model_optim.step() | |||
| return metric,loss.item() | |||
| def _train_controller(self, epoch): | |||
| self.model.eval() | |||
| self.controller.train() | |||
| self.ctrl_optim.zero_grad() | |||
| rewards=[] | |||
| for ctrl_step in range(self.ctrl_steps_aggregate): | |||
| self._resample() | |||
| with torch.no_grad(): | |||
| metric,loss=self._infer() | |||
| reward =-metric # todo : now metric is loss | |||
| metric,loss=self._infer(mask='val') | |||
| reward =metric | |||
| rewards.append(reward) | |||
| if self.entropy_weight: | |||
| reward += self.entropy_weight * self.controller.sample_entropy.item() | |||
| self.baseline = self.baseline * self.baseline_decay + reward * (1 - self.baseline_decay) | |||
| @@ -357,6 +174,7 @@ class Enas(BaseNAS): | |||
| if self.log_frequency is not None and ctrl_step % self.log_frequency == 0: | |||
| _logger.info('RL Epoch [%d/%d] Step [%d/%d] %s', epoch + 1, self.num_epochs, | |||
| ctrl_step + 1, self.ctrl_steps_aggregate) | |||
| return sum(rewards)/len(rewards) | |||
| def _resample(self): | |||
| result = self.controller.resample() | |||
| @@ -368,6 +186,6 @@ class Enas(BaseNAS): | |||
| with torch.no_grad(): | |||
| return self.controller.resample() | |||
| def _infer(self): | |||
| metric, loss = self.estimator.infer(self.model, self.dataset) | |||
| def _infer(self,mask='train'): | |||
| metric, loss = self.estimator.infer(self.model, self.dataset,mask=mask) | |||
| return metric, loss | |||
+ 383
- 0
autogl/module/nas/algorithm/rl.py
View File
| @@ -0,0 +1,383 @@ | |||
| # codes in this file are reproduced from https://github.com/microsoft/nni with some changes. | |||
| import copy | |||
| import logging | |||
| import torch | |||
| import torch.nn as nn | |||
| import torch.nn.functional as F | |||
| from .base import BaseNAS | |||
| from ..space import BaseSpace | |||
| from ..utils import AverageMeterGroup, replace_layer_choice, replace_input_choice, get_module_order, sort_replaced_module | |||
| from nni.nas.pytorch.fixed import apply_fixed_architecture | |||
| from tqdm import tqdm | |||
| _logger = logging.getLogger(__name__) | |||
| def _get_mask(sampled, total): | |||
| multihot = [i == sampled or (isinstance(sampled, list) and i in sampled) for i in range(total)] | |||
| return torch.tensor(multihot, dtype=torch.bool) # pylint: disable=not-callable | |||
| class PathSamplingLayerChoice(nn.Module): | |||
| """ | |||
| Mixed module, in which fprop is decided by exactly one or multiple (sampled) module. | |||
| If multiple module is selected, the result will be sumed and returned. | |||
| Attributes | |||
| ---------- | |||
| sampled : int or list of int | |||
| Sampled module indices. | |||
| mask : tensor | |||
| A multi-hot bool 1D-tensor representing the sampled mask. | |||
| """ | |||
| def __init__(self, layer_choice): | |||
| super(PathSamplingLayerChoice, self).__init__() | |||
| self.op_names = [] | |||
| for name, module in layer_choice.named_children(): | |||
| self.add_module(name, module) | |||
| self.op_names.append(name) | |||
| assert self.op_names, 'There has to be at least one op to choose from.' | |||
| self.sampled = None # sampled can be either a list of indices or an index | |||
| def forward(self, *args, **kwargs): | |||
| assert self.sampled is not None, 'At least one path needs to be sampled before fprop.' | |||
| if isinstance(self.sampled, list): | |||
| return sum([getattr(self, self.op_names[i])(*args, **kwargs) for i in self.sampled]) # pylint: disable=not-an-iterable | |||
| else: | |||
| return getattr(self, self.op_names[self.sampled])(*args, **kwargs) # pylint: disable=invalid-sequence-index | |||
| def __len__(self): | |||
| return len(self.op_names) | |||
| @property | |||
| def mask(self): | |||
| return _get_mask(self.sampled, len(self)) | |||
| class PathSamplingInputChoice(nn.Module): | |||
| """ | |||
| Mixed input. Take a list of tensor as input, select some of them and return the sum. | |||
| Attributes | |||
| ---------- | |||
| sampled : int or list of int | |||
| Sampled module indices. | |||
| mask : tensor | |||
| A multi-hot bool 1D-tensor representing the sampled mask. | |||
| """ | |||
| def __init__(self, input_choice): | |||
| super(PathSamplingInputChoice, self).__init__() | |||
| self.n_candidates = input_choice.n_candidates | |||
| self.n_chosen = input_choice.n_chosen | |||
| self.sampled = None | |||
| def forward(self, input_tensors): | |||
| if isinstance(self.sampled, list): | |||
| return sum([input_tensors[t] for t in self.sampled]) # pylint: disable=not-an-iterable | |||
| else: | |||
| return input_tensors[self.sampled] | |||
| def __len__(self): | |||
| return self.n_candidates | |||
| @property | |||
| def mask(self): | |||
| return _get_mask(self.sampled, len(self)) | |||
| def __repr__(self): | |||
| return f'PathSamplingInputChoice(n_candidates={self.n_candidates}, chosen={self.sampled})' | |||
| class StackedLSTMCell(nn.Module): | |||
| def __init__(self, layers, size, bias): | |||
| super().__init__() | |||
| self.lstm_num_layers = layers | |||
| self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias) | |||
| for _ in range(self.lstm_num_layers)]) | |||
| def forward(self, inputs, hidden): | |||
| prev_h, prev_c = hidden | |||
| next_h, next_c = [], [] | |||
| for i, m in enumerate(self.lstm_modules): | |||
| curr_h, curr_c = m(inputs, (prev_h[i], prev_c[i])) | |||
| next_c.append(curr_c) | |||
| next_h.append(curr_h) | |||
| # current implementation only supports batch size equals 1, | |||
| # but the algorithm does not necessarily have this limitation | |||
| inputs = curr_h[-1].view(1, -1) | |||
| return next_h, next_c | |||
| class ReinforceField: | |||
| """ | |||
| A field with ``name``, with ``total`` choices. ``choose_one`` is true if one and only one is meant to be | |||
| selected. Otherwise, any number of choices can be chosen. | |||
| """ | |||
| def __init__(self, name, total, choose_one): | |||
| self.name = name | |||
| self.total = total | |||
| self.choose_one = choose_one | |||
| def __repr__(self): | |||
| return f'ReinforceField(name={self.name}, total={self.total}, choose_one={self.choose_one})' | |||
| class ReinforceController(nn.Module): | |||
| """ | |||
| A controller that mutates the graph with RL. | |||
| Parameters | |||
| ---------- | |||
| fields : list of ReinforceField | |||
| List of fields to choose. | |||
| lstm_size : int | |||
| Controller LSTM hidden units. | |||
| lstm_num_layers : int | |||
| Number of layers for stacked LSTM. | |||
| tanh_constant : float | |||
| Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``. | |||
| skip_target : float | |||
| Target probability that skipconnect will appear. | |||
| temperature : float | |||
| Temperature constant that divides the logits. | |||
| entropy_reduction : str | |||
| Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced. | |||
| """ | |||
| def __init__(self, fields, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, | |||
| skip_target=0.4, temperature=None, entropy_reduction='sum'): | |||
| super(ReinforceController, self).__init__() | |||
| self.fields = fields | |||
| self.lstm_size = lstm_size | |||
| self.lstm_num_layers = lstm_num_layers | |||
| self.tanh_constant = tanh_constant | |||
| self.temperature = temperature | |||
| self.skip_target = skip_target | |||
| self.lstm = StackedLSTMCell(self.lstm_num_layers, self.lstm_size, False) | |||
| self.attn_anchor = nn.Linear(self.lstm_size, self.lstm_size, bias=False) | |||
| self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False) | |||
| self.v_attn = nn.Linear(self.lstm_size, 1, bias=False) | |||
| self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1) | |||
| self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), # pylint: disable=not-callable | |||
| requires_grad=False) | |||
| assert entropy_reduction in ['sum', 'mean'], 'Entropy reduction must be one of sum and mean.' | |||
| self.entropy_reduction = torch.sum if entropy_reduction == 'sum' else torch.mean | |||
| self.cross_entropy_loss = nn.CrossEntropyLoss(reduction='none') | |||
| self.soft = nn.ModuleDict({ | |||
| field.name: nn.Linear(self.lstm_size, field.total, bias=False) for field in fields | |||
| }) | |||
| self.embedding = nn.ModuleDict({ | |||
| field.name: nn.Embedding(field.total, self.lstm_size) for field in fields | |||
| }) | |||
| def resample(self): | |||
| self._initialize() | |||
| result = dict() | |||
| for field in self.fields: | |||
| result[field.name] = self._sample_single(field) | |||
| return result | |||
| def _initialize(self): | |||
| self._inputs = self.g_emb.data | |||
| self._c = [torch.zeros((1, self.lstm_size), | |||
| dtype=self._inputs.dtype, | |||
| device=self._inputs.device) for _ in range(self.lstm_num_layers)] | |||
| self._h = [torch.zeros((1, self.lstm_size), | |||
| dtype=self._inputs.dtype, | |||
| device=self._inputs.device) for _ in range(self.lstm_num_layers)] | |||
| self.sample_log_prob = 0 | |||
| self.sample_entropy = 0 | |||
| self.sample_skip_penalty = 0 | |||
| def _lstm_next_step(self): | |||
| self._h, self._c = self.lstm(self._inputs, (self._h, self._c)) | |||
| def _sample_single(self, field): | |||
| self._lstm_next_step() | |||
| logit = self.soft[field.name](self._h[-1]) | |||
| if self.temperature is not None: | |||
| logit /= self.temperature | |||
| if self.tanh_constant is not None: | |||
| logit = self.tanh_constant * torch.tanh(logit) | |||
| if field.choose_one: | |||
| sampled = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1) | |||
| log_prob = self.cross_entropy_loss(logit, sampled) | |||
| self._inputs = self.embedding[field.name](sampled) | |||
| else: | |||
| logit = logit.view(-1, 1) | |||
| logit = torch.cat([-logit, logit], 1) # pylint: disable=invalid-unary-operand-type | |||
| sampled = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1) | |||
| skip_prob = torch.sigmoid(logit) | |||
| kl = torch.sum(skip_prob * torch.log(skip_prob / self.skip_targets)) | |||
| self.sample_skip_penalty += kl | |||
| log_prob = self.cross_entropy_loss(logit, sampled) | |||
| sampled = sampled.nonzero().view(-1) | |||
| if sampled.sum().item(): | |||
| self._inputs = (torch.sum(self.embedding[field.name](sampled.view(-1)), 0) / (1. + torch.sum(sampled))).unsqueeze(0) | |||
| else: | |||
| self._inputs = torch.zeros(1, self.lstm_size, device=self.embedding[field.name].weight.device) | |||
| sampled = sampled.detach().numpy().tolist() | |||
| self.sample_log_prob += self.entropy_reduction(log_prob) | |||
| entropy = (log_prob * torch.exp(-log_prob)).detach() # pylint: disable=invalid-unary-operand-type | |||
| self.sample_entropy += self.entropy_reduction(entropy) | |||
| if len(sampled) == 1: | |||
| sampled = sampled[0] | |||
| return sampled | |||
| class RL(BaseNAS): | |||
| """ | |||
| ENAS trainer. | |||
| Parameters | |||
| ---------- | |||
| model : nn.Module | |||
| PyTorch model to be trained. | |||
| loss : callable | |||
| Receives logits and ground truth label, return a loss tensor. | |||
| metrics : callable | |||
| Receives logits and ground truth label, return a dict of metrics. | |||
| reward_function : callable | |||
| Receives logits and ground truth label, return a tensor, which will be feeded to RL controller as reward. | |||
| optimizer : Optimizer | |||
| The optimizer used for optimizing the model. | |||
| num_epochs : int | |||
| Number of epochs planned for training. | |||
| dataset : Dataset | |||
| Dataset for training. Will be split for training weights and architecture weights. | |||
| batch_size : int | |||
| Batch size. | |||
| workers : int | |||
| Workers for data loading. | |||
| device : torch.device | |||
| ``torch.device("cpu")`` or ``torch.device("cuda")``. | |||
| log_frequency : int | |||
| Step count per logging. | |||
| grad_clip : float | |||
| Gradient clipping. Set to 0 to disable. Default: 5. | |||
| entropy_weight : float | |||
| Weight of sample entropy loss. | |||
| skip_weight : float | |||
| Weight of skip penalty loss. | |||
| baseline_decay : float | |||
| Decay factor of baseline. New baseline will be equal to ``baseline_decay * baseline_old + reward * (1 - baseline_decay)``. | |||
| ctrl_lr : float | |||
| Learning rate for RL controller. | |||
| ctrl_steps_aggregate : int | |||
| Number of steps that will be aggregated into one mini-batch for RL controller. | |||
| ctrl_steps : int | |||
| Number of mini-batches for each epoch of RL controller learning. | |||
| ctrl_kwargs : dict | |||
| Optional kwargs that will be passed to :class:`ReinforceController`. | |||
| """ | |||
| def __init__(self, device='cuda', workers=4,log_frequency=None, | |||
| grad_clip=5., entropy_weight=0.0001, skip_weight=0.8, baseline_decay=0.999, | |||
| ctrl_lr=0.00035, ctrl_steps_aggregate=20, ctrl_kwargs=None,n_warmup=100,model_lr=5e-3,model_wd=5e-4,*args,**kwargs): | |||
| super().__init__(device) | |||
| self.device=device | |||
| self.num_epochs = kwargs.get("num_epochs", 5) | |||
| self.workers = workers | |||
| self.log_frequency = log_frequency | |||
| self.entropy_weight = entropy_weight | |||
| self.skip_weight = skip_weight | |||
| self.baseline_decay = baseline_decay | |||
| self.baseline = 0. | |||
| self.ctrl_steps_aggregate = ctrl_steps_aggregate | |||
| self.grad_clip = grad_clip | |||
| self.workers = workers | |||
| self.ctrl_kwargs=ctrl_kwargs | |||
| self.ctrl_lr=ctrl_lr | |||
| self.n_warmup=n_warmup | |||
| self.model_lr = model_lr | |||
| self.model_wd = model_wd | |||
| self.log=open('log.txt','w') | |||
| 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] | |||
| self.controller = ReinforceController(self.nas_fields, **(self.ctrl_kwargs or {})) | |||
| self.ctrl_optim = torch.optim.Adam(self.controller.parameters(), lr=self.ctrl_lr) | |||
| # train | |||
| with tqdm(range(self.num_epochs)) as bar: | |||
| for i in bar: | |||
| l2=self._train_controller(i) | |||
| # try: | |||
| # l2=self._train_controller(i) | |||
| # except Exception as e: | |||
| # print(e) | |||
| # nm=self.nas_modules | |||
| # for i in range(len(nm)): | |||
| # print(nm[i][1].sampled) | |||
| # # import pdb | |||
| # # pdb.set_trace() | |||
| bar.set_postfix(reward_controller=l2) | |||
| selection=self.export() | |||
| arch=space.export(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=[] | |||
| with tqdm(range(self.ctrl_steps_aggregate)) as bar: | |||
| for ctrl_step in bar: | |||
| self._resample() | |||
| metric,loss=self._infer(mask='val') | |||
| bar.set_postfix(acc=metric,loss=loss.item()) | |||
| self.log.write(f'{self.arch}\n{self.selection}\n{metric},{loss}\n') | |||
| self.log.flush() | |||
| reward =metric | |||
| rewards.append(reward) | |||
| if self.entropy_weight: | |||
| reward += self.entropy_weight * self.controller.sample_entropy.item() | |||
| self.baseline = self.baseline * self.baseline_decay + reward * (1 - self.baseline_decay) | |||
| loss = self.controller.sample_log_prob * (reward - self.baseline) | |||
| if self.skip_weight: | |||
| loss += self.skip_weight * self.controller.sample_skip_penalty | |||
| loss /= self.ctrl_steps_aggregate | |||
| loss.backward() | |||
| if (ctrl_step + 1) % self.ctrl_steps_aggregate == 0: | |||
| if self.grad_clip > 0: | |||
| nn.utils.clip_grad_norm_(self.controller.parameters(), self.grad_clip) | |||
| self.ctrl_optim.step() | |||
| self.ctrl_optim.zero_grad() | |||
| if self.log_frequency is not None and ctrl_step % self.log_frequency == 0: | |||
| _logger.info('RL Epoch [%d/%d] Step [%d/%d] %s', epoch + 1, self.num_epochs, | |||
| ctrl_step + 1, self.ctrl_steps_aggregate) | |||
| return sum(rewards)/len(rewards) | |||
| def _resample(self): | |||
| result = self.controller.resample() | |||
| self.arch=self.model.export(result,device=self.device) | |||
| self.selection=result | |||
| def export(self): | |||
| self.controller.eval() | |||
| with torch.no_grad(): | |||
| return self.controller.resample() | |||
| def _infer(self,mask='train'): | |||
| metric, loss = self.estimator.infer(self.arch, self.dataset,mask=mask) | |||
| return metric, loss | |||
+ 26
- 2
autogl/module/nas/estimator/one_shot.py
View File
| @@ -3,7 +3,7 @@ import torch.nn.functional as F | |||
| from ..space import BaseSpace | |||
| from .base import BaseEstimator | |||
| import torch | |||
| class OneShotEstimator(BaseEstimator): | |||
| """ | |||
| @@ -18,4 +18,28 @@ class OneShotEstimator(BaseEstimator): | |||
| pred = model(dset)[getattr(dset, f"{mask}_mask")] | |||
| y = dset.y[getattr(dset, f'{mask}_mask')] | |||
| loss = F.nll_loss(pred, y) | |||
| return loss, loss | |||
| acc=sum(pred.max(1)[1]==y).item()/y.size(0) | |||
| return acc, loss | |||
| from autogl.module.train import NodeClassificationFullTrainer | |||
| class TrainEstimator(BaseEstimator): | |||
| def __init__(self): | |||
| self.estimator=OneShotEstimator() | |||
| def infer(self,model: BaseSpace, dataset, mask="train"): | |||
| # self.trainer.model=model | |||
| # self.trainer.device=model.device | |||
| self.trainer=NodeClassificationFullTrainer( | |||
| model=model, | |||
| optimizer=torch.optim.Adam, | |||
| lr=0.01, | |||
| max_epoch=200, | |||
| early_stopping_round=200, | |||
| weight_decay=5e-4, | |||
| device="auto", | |||
| init=False, | |||
| feval=['acc'], | |||
| loss="nll_loss", | |||
| lr_scheduler_type=None) | |||
| self.trainer.train(dataset) | |||
| with torch.no_grad(): | |||
| return self.estimator.infer(model,dataset,mask) | |||
+ 29
- 4
autogl/module/nas/space/graph_nas.py
View File
| @@ -43,6 +43,9 @@ class LambdaModule(nn.Module): | |||
| def forward(self, x): | |||
| return self.lambd(x) | |||
| def __repr__(self): | |||
| return '{}({})'.format(self.__class__.__name__,self.lambd) | |||
| class StrModule(nn.Module): | |||
| def __init__(self, lambd): | |||
| super().__init__() | |||
| @@ -50,6 +53,9 @@ class StrModule(nn.Module): | |||
| def forward(self, *args,**kwargs): | |||
| return self.str | |||
| def __repr__(self): | |||
| return '{}({})'.format(self.__class__.__name__,self.str) | |||
| def act_map(act): | |||
| if act == "linear": | |||
| return lambda x: x | |||
| @@ -108,7 +114,11 @@ def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> nn.Module: | |||
| elif gnn_name == "linear": | |||
| return LinearConv(in_dim, out_dim, bias=bias) | |||
| elif gnn_name == "zero": | |||
| return ZeroConv(in_dim, out_dim, bias=bias) | |||
| # return ZeroConv(in_dim, out_dim, bias=bias) | |||
| return Identity() | |||
| class Identity(nn.Module): | |||
| def forward(self, x, edge_index, edge_weight=None): | |||
| return x | |||
| class LinearConv(nn.Module): | |||
| def __init__(self, | |||
| in_channels, | |||
| @@ -128,6 +138,15 @@ class LinearConv(nn.Module): | |||
| self.out_channels) | |||
| from torch.autograd import Function | |||
| class ZeroConvFunc(Function): | |||
| @staticmethod | |||
| def forward(ctx,x): | |||
| return x | |||
| @staticmethod | |||
| def backward(ctx, grad_output): | |||
| return 0 | |||
| class ZeroConv(nn.Module): | |||
| def __init__(self, | |||
| in_channels, | |||
| @@ -138,9 +157,8 @@ class ZeroConv(nn.Module): | |||
| self.out_channels = out_channels | |||
| self.out_dim = out_channels | |||
| def forward(self, x, edge_index, edge_weight=None): | |||
| return torch.zeros([x.size(0), self.out_dim]).to(x.device) | |||
| return ZeroConvFunc.apply(torch.zeros([x.size(0), self.out_dim]).to(x.device)) | |||
| def __repr__(self): | |||
| return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, | |||
| @@ -193,6 +211,8 @@ class GraphNasNodeClassificationSpace(BaseSpace): | |||
| setattr(self,f"act",self.setLayerChoice(2*layer,[act_map_nn(a)for a in act_list],key=f"act")) | |||
| setattr(self,f"concat",self.setLayerChoice(2*layer+1,map_nn(["add", "product", "concat"]) ,key=f"concat")) | |||
| self._initialized = True | |||
| self.classifier1 = nn.Linear(self.hidden_dim*self.layer_number, self.output_dim) | |||
| self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim) | |||
| def forward(self, data): | |||
| x, edges = data.x, data.edge_index # x [2708,1433] ,[2, 10556] | |||
| @@ -202,10 +222,11 @@ class GraphNasNodeClassificationSpace(BaseSpace): | |||
| node_in = getattr(self, f"in_{layer}")(prev_nodes_out) | |||
| node_out= getattr(self, f"op_{layer}")(node_in,edges) | |||
| prev_nodes_out.append(node_out) | |||
| if self.search_act_con: | |||
| if not self.search_act_con: | |||
| x = torch.cat(prev_nodes_out[2:],dim=1) | |||
| x = F.leaky_relu(x) | |||
| x = F.dropout(x, p=self.dropout, training = self.training) | |||
| x = self.classifier1(x) | |||
| else: | |||
| act=getattr(self, f"act") | |||
| con=getattr(self, f"concat")() | |||
| @@ -222,6 +243,10 @@ class GraphNasNodeClassificationSpace(BaseSpace): | |||
| x=tmp | |||
| x = act(x) | |||
| x = F.dropout(x, p=self.dropout, training = self.training) | |||
| if con=='concat': | |||
| x=self.classifier1(x) | |||
| else: | |||
| x=self.classifier2(x) | |||
| return F.log_softmax(x, dim=1) | |||
| def export(self, selection, device) -> BaseModel: | |||
+ 2
- 2
examples/test_graph_nas.py
View File
| @@ -28,9 +28,9 @@ if __name__ == '__main__': | |||
| feval=['acc'], | |||
| loss="nll_loss", | |||
| lr_scheduler_type=None,), | |||
| nas_algorithms=[Enas(num_epochs=10)], | |||
| nas_algorithms=[Enas(num_epochs=400,n_warmup=250)], | |||
| #nas_algorithms=[Darts(num_epochs=200)], | |||
| nas_spaces=[GraphNasNodeClassificationSpace(hidden_dim=16, ops=[GCNConv, GCNConv],search_act_con=True)], | |||
| nas_spaces=[GraphNasNodeClassificationSpace(hidden_dim=32,search_act_con=False,layer_number=2)], | |||
| nas_estimators=[OneShotEstimator()] | |||
| ) | |||
| solver.fit(dataset) | |||
+ 42
- 0
examples/test_graph_nas_rl.py
View File
| @@ -0,0 +1,42 @@ | |||
| import sys | |||
| sys.path.append('../') | |||
| from torch_geometric.nn import GCNConv | |||
| import torch | |||
| from autogl.datasets import build_dataset_from_name | |||
| from autogl.solver import AutoNodeClassifier | |||
| from autogl.module.train import NodeClassificationFullTrainer | |||
| from autogl.module.nas import Darts, OneShotEstimator | |||
| from autogl.module.nas.space.graph_nas import GraphNasNodeClassificationSpace | |||
| from autogl.module.train import Acc | |||
| from autogl.module.nas.algorithm.enas import Enas | |||
| from autogl.module.nas.algorithm.rl import RL | |||
| from autogl.module.nas.estimator.one_shot import TrainEstimator | |||
| import logging | |||
| if __name__ == '__main__': | |||
| logging.getLogger().setLevel(logging.WARNING) | |||
| dataset = build_dataset_from_name('cora') | |||
| solver = AutoNodeClassifier( | |||
| feature_module='PYGNormalizeFeatures', | |||
| graph_models=[], | |||
| hpo_module=None, | |||
| ensemble_module=None, | |||
| default_trainer=NodeClassificationFullTrainer( | |||
| optimizer=torch.optim.Adam, | |||
| lr=0.01, | |||
| max_epoch=200, | |||
| early_stopping_round=200, | |||
| weight_decay=5e-4, | |||
| device="auto", | |||
| init=False, | |||
| feval=['acc'], | |||
| loss="nll_loss", | |||
| lr_scheduler_type=None,), | |||
| nas_algorithms=[RL(num_epochs=400)], | |||
| #nas_algorithms=[Darts(num_epochs=200)], | |||
| nas_spaces=[GraphNasNodeClassificationSpace(hidden_dim=16,search_act_con=True,layer_number=2)], | |||
| nas_estimators=[TrainEstimator()] | |||
| ) | |||
| solver.fit(dataset) | |||
| solver.get_leaderboard().show() | |||
| out = solver.predict_proba() | |||
| print('acc on cora', Acc.evaluate(out, dataset[0].y[dataset[0].test_mask].detach().numpy())) | |||
+ 73
- 0
examples/test_graph_nas_space.py
View File
| @@ -0,0 +1,73 @@ | |||
| import sys | |||
| from nni.nas.pytorch.mutables import Mutable | |||
| sys.path.append('../') | |||
| from torch_geometric.nn import GCNConv | |||
| import torch | |||
| from autogl.datasets import build_dataset_from_name | |||
| from autogl.solver import AutoNodeClassifier | |||
| from autogl.module.train import NodeClassificationFullTrainer | |||
| from autogl.module.nas import Darts, OneShotEstimator | |||
| from autogl.module.nas.space.graph_nas import * | |||
| from autogl.module.train import Acc | |||
| from autogl.module.nas.algorithm.enas import Enas | |||
| from autogl.module.nas.algorithm.rl import * | |||
| from autogl.module.nas.estimator.one_shot import TrainEstimator | |||
| import logging | |||
| import numpy as np | |||
| from tqdm import tqdm | |||
| if __name__ == '__main__': | |||
| logging.getLogger().setLevel(logging.WARNING) | |||
| dataset = build_dataset_from_name('cora') | |||
| space=GraphNasNodeClassificationSpace(hidden_dim=16,search_act_con=True,layer_number=2) | |||
| space.instantiate(input_dim=dataset[0].x.shape[1], | |||
| output_dim=dataset.num_classes,) | |||
| estim=TrainEstimator() | |||
| # solver.fit(dataset) | |||
| # solver.get_leaderboard().show() | |||
| # out = solver.predict_proba() | |||
| # print('acc on cora', Acc.evaluate(out, dataset[0].y[dataset[0].test_mask].detach().numpy())) | |||
| class Tmp: | |||
| def __init__(self,space): | |||
| self.model = space | |||
| 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) | |||
| t=Tmp(space) | |||
| print(t.nas_modules) | |||
| nm=t.nas_modules | |||
| selection_range={} | |||
| for k,v in nm: | |||
| selection_range[k]=len(v) | |||
| ks=list(selection_range.keys()) | |||
| selections=[] | |||
| def dfs(selection,d): | |||
| if d>=len(ks): | |||
| selections.append(selection.copy()) | |||
| return | |||
| k=ks[d] | |||
| r=selection_range[k] | |||
| for i in range(r): | |||
| selection[k]=i | |||
| dfs(selection,d+1) | |||
| dfs({},0) | |||
| print(f'#selections {len(selections)}') | |||
| device=torch.device('cuda:0') | |||
| accs=[] | |||
| from datetime import datetime | |||
| timestamp=datetime.now().strftime('%m%d-%H-%M-%S') | |||
| log=open(f'acclog{timestamp}.txt','w') | |||
| with tqdm(selections) as bar: | |||
| for selection in bar: | |||
| arch=space.export(selection,device) | |||
| m,l=estim.infer(arch,dataset,'test') | |||
| bar.set_postfix(m=m,l=l.item()) | |||
| log.write(f'{arch}\n{selection}\n{m},{l}\n') | |||
| log.flush() | |||
| accs.append(m) | |||
| np.save(f'space_acc{timestamp}',np.array(accs)) | |||
| print(f'max acc {np.max(accs)}') | |||