Merge pull request #27 from THUMNLab/graphnas

Graphnas

autogl/module/nas/algorithm/enas.py

@@ -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

autogl/module/nas/algorithm/rl.py

@@ -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

autogl/module/nas/estimator/one_shot.py

@@ -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)

autogl/module/nas/space/graph_nas.py

@@ -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:

examples/test_graph_nas.py

@@ -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)

examples/test_graph_nas_rl.py

@@ -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()))

examples/test_graph_nas_space.py

@@ -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)}')