From 62ee4808e6b4032f9fc8df3abeae8f41370f0be8 Mon Sep 17 00:00:00 2001
From: cluster32 <general502570@outlook.com>
Date: Wed, 24 Mar 2021 11:37:30 +0800
Subject: [PATCH] test nas

---
 autogl/module/hpo/__init__.py   |   2 +
 autogl/module/hpo/darts.py      | 184 ++++++++++++++++++++++++++++++++
 autogl/module/hpo/nas.py        |  46 ++++++++
 autogl/module/hpo/test.py       |  33 ++++++
 autogl/module/hpo/utils.py      | 182 +++++++++++++++++++++++++++++++
 examples/node_classification.py |   2 +-
 6 files changed, 448 insertions(+), 1 deletion(-)
 create mode 100644 autogl/module/hpo/darts.py
 create mode 100644 autogl/module/hpo/nas.py
 create mode 100644 autogl/module/hpo/test.py
 create mode 100644 autogl/module/hpo/utils.py

diff --git a/autogl/module/hpo/__init__.py b/autogl/module/hpo/__init__.py
index e8fe41a..cd4bbba 100644
--- a/autogl/module/hpo/__init__.py
+++ b/autogl/module/hpo/__init__.py
@@ -28,6 +28,7 @@ from .mocmaes_advisorchoco import MocmaesAdvisorChoco
 from .quasi_advisorchoco import QuasiAdvisorChoco
 from .rand_advisor import RandAdvisor
 from .tpe_advisorhpo import TpeAdvisorHPO
+from .test import TestHPO
 
 
 def build_hpo_from_name(name: str) -> BaseHPOptimizer:
@@ -62,5 +63,6 @@ __all__ = [
     "QuasiAdvisorChoco",
     "RandAdvisor",
     "TpeAdvisorHPO",
+    "TestHPO",
     "build_hpo_from_name",
 ]
diff --git a/autogl/module/hpo/darts.py b/autogl/module/hpo/darts.py
new file mode 100644
index 0000000..6ac4b86
--- /dev/null
+++ b/autogl/module/hpo/darts.py
@@ -0,0 +1,184 @@
+# Modified from NNI
+
+import copy
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .nas import BaseNAS
+from .utils import AverageMeterGroup, replace_layer_choice, replace_input_choice
+
+
+_logger = logging.getLogger(__name__)
+
+
+class DartsLayerChoice(nn.Module):
+    def __init__(self, layer_choice):
+        super(DartsLayerChoice, self).__init__()
+        self.name = layer_choice.key
+        self.op_choices = nn.ModuleDict(layer_choice.named_children())
+        self.alpha = nn.Parameter(torch.randn(len(self.op_choices)) * 1e-3)
+
+    def forward(self, *args, **kwargs):
+        op_results = torch.stack([op(*args, **kwargs) for op in self.op_choices.values()])
+        alpha_shape = [-1] + [1] * (len(op_results.size()) - 1)
+        return torch.sum(op_results * F.softmax(self.alpha, -1).view(*alpha_shape), 0)
+
+    def parameters(self):
+        for _, p in self.named_parameters():
+            yield p
+
+    def named_parameters(self):
+        for name, p in super(DartsLayerChoice, self).named_parameters():
+            if name == 'alpha':
+                continue
+            yield name, p
+
+    def export(self):
+        return torch.argmax(self.alpha).item()
+
+
+class DartsInputChoice(nn.Module):
+    def __init__(self, input_choice):
+        super(DartsInputChoice, self).__init__()
+        self.name = input_choice.key
+        self.alpha = nn.Parameter(torch.randn(input_choice.n_candidates) * 1e-3)
+        self.n_chosen = input_choice.n_chosen or 1
+
+    def forward(self, inputs):
+        inputs = torch.stack(inputs)
+        alpha_shape = [-1] + [1] * (len(inputs.size()) - 1)
+        return torch.sum(inputs * F.softmax(self.alpha, -1).view(*alpha_shape), 0)
+
+    def parameters(self):
+        for _, p in self.named_parameters():
+            yield p
+
+    def named_parameters(self):
+        for name, p in super(DartsInputChoice, self).named_parameters():
+            if name == 'alpha':
+                continue
+            yield name, p
+
+    def export(self):
+        return torch.argsort(-self.alpha).cpu().numpy().tolist()[:self.n_chosen]
+
+
+class DartsTrainer(BaseNAS):
+    """
+    DARTS 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.
+    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.
+    grad_clip : float
+        Gradient clipping. Set to 0 to disable. Default: 5.
+    learning_rate : float
+        Learning rate to optimize the model.
+    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.
+    arc_learning_rate : float
+        Learning rate of architecture parameters.
+    unrolled : float
+        ``True`` if using second order optimization, else first order optimization.
+    """
+
+    """def __init__(self, model, loss, metrics, optimizer,
+                 num_epochs, dataset, grad_clip=5.,
+                 learning_rate=2.5E-3, batch_size=64, workers=4,
+                 device=None, log_frequency=None,
+                 arc_learning_rate=3.0E-4, unrolled=False):"""
+    def __init__(self, *args, **kwargs):
+        self.num_epochs = kwargs.get("num_epochs", 5)
+        self.workers = 4
+        self.device = "cuda"
+        self.log_frequency = None
+
+        #for _, module in self.nas_modules:
+        #    module.to(self.device)
+
+        # use the same architecture weight for modules with duplicated names
+
+    def search(self, space, dset, trainer):
+        """
+        main process
+        """
+        self.model = space
+        self.dataset = dset
+        self.trainer = trainer
+        self.model_optim = torch.optim.SGD(
+            self.model.parameters(), lr=0.01, weight_decay=3e-4
+        )
+
+        self.nas_modules = []
+        replace_layer_choice(self.model, DartsLayerChoice, self.nas_modules)
+        replace_input_choice(self.model, DartsInputChoice, self.nas_modules)
+
+        ctrl_params = {}
+        for _, m in self.nas_modules:
+            if m.name in ctrl_params:
+                assert m.alpha.size() == ctrl_params[m.name].size(), 'Size of parameters with the same label should be same.'
+                m.alpha = ctrl_params[m.name]
+            else:
+                ctrl_params[m.name] = m.alpha
+        self.ctrl_optim = torch.optim.Adam(list(ctrl_params.values()), 3e-4, betas=(0.5, 0.999),
+                                           weight_decay=1.0E-3)
+        self.grad_clip = 5.
+
+        for step in range(self.num_epochs):
+            self._train_one_epoch(step)
+            if self.log_frequency is not None and step % self.log_frequency == 0:
+                _logger.info('Epoch [%s/%s] Step [%s/%s]  %s', epoch + 1,
+                                self.num_epochs, step + 1, len(self.train_loader), meters)
+
+        return self.export()
+
+    def _train_one_epoch(self, epoch):
+        self.model.train()
+        meters = AverageMeterGroup()
+
+        # phase 1. architecture step
+        self.ctrl_optim.zero_grad()
+        # only no unroll here
+        _, loss = self._infer()
+        loss.backward()
+        self.ctrl_optim.step()
+
+        # phase 2: child network step
+        self.model_optim.zero_grad()
+        metric, loss = self._infer()
+        loss.backward()
+        if self.grad_clip > 0:
+            nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)  # gradient clipping
+        self.model_optim.step()
+
+    def _infer(self):
+        metric, loss = self.trainer.infer(self.model, self.dataset)
+        return metric, loss
+
+    @torch.no_grad()
+    def export(self):
+        result = dict()
+        for name, module in self.nas_modules:
+            if name not in result:
+                result[name] = module.export()
+        return result
diff --git a/autogl/module/hpo/nas.py b/autogl/module/hpo/nas.py
new file mode 100644
index 0000000..8091d06
--- /dev/null
+++ b/autogl/module/hpo/nas.py
@@ -0,0 +1,46 @@
+from torch_geometric.nn import GCNConv, SAGEConv
+from nni.nas.pytorch import mutables
+import torch.nn as nn
+
+class BaseNAS:
+    def search(self, space, dset, trainer):
+        """
+        The main process of NAS.
+        Parameters
+        ----------
+        space : BaseArchitectureSpace
+            No implementation yet
+        dataset : ...datasets
+            Dataset to train and evaluate.
+        trainer : ..train.BaseTrainer
+            Including model, giving HP space and using for training
+
+        Returns
+        -------
+        model: ..train.BaseTrainer
+            The trainer including the best trained model
+        """
+
+class GraphSpace(nn.Module):
+    def __init__(self, inp, hid, oup):
+        super().__init__()
+        self.gcn = GCNConv(inp, hid)
+        self.op1 = mutables.LayerChoice([GCNConv(inp, hid),SAGEConv(inp, hid)])
+        self.op2 = mutables.LayerChoice([
+            GCNConv(hid, oup),
+            SAGEConv(hid, oup)       
+        ], key = "2")
+
+    def forward(self, data):
+        x = self.op1(data.x, data.edge_index)
+        x = self.op2(x, data.edge_index)
+        return x
+        
+class BaseTrainer:
+    def infer(self, model, dataset):
+        dset = dataset[0]
+        pred = model(dset)[dset.train_mask]
+        y = dset.y[dset.train_mask]
+        loss_func = nn.CrossEntropyLoss()
+        loss = loss_func(pred, y)
+        return loss, loss
\ No newline at end of file
diff --git a/autogl/module/hpo/test.py b/autogl/module/hpo/test.py
new file mode 100644
index 0000000..fdb0a33
--- /dev/null
+++ b/autogl/module/hpo/test.py
@@ -0,0 +1,33 @@
+import hyperopt
+
+from . import register_hpo
+from .nas import BaseTrainer, GraphSpace
+from .darts import DartsTrainer
+from .base import BaseHPOptimizer, TimeTooLimitedError
+
+@register_hpo("test")
+class TestHPO(BaseHPOptimizer):
+    """
+    Random search algorithm in `advisor` package
+    See https://github.com/tobegit3hub/advisor for the package
+    See .advisorbase.AdvisorBaseHPOptimizer for more information
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def optimize(self, trainer, dataset, time_limit=None, memory_limit=None):
+        num_features=dataset[0].x.shape[1]
+        num_classes=dataset.num_classes
+        model = GraphSpace(num_features, 64, num_classes)
+        tr = BaseTrainer()
+        nas = DartsTrainer()
+        a = nas.search(model, dataset, tr)
+        print(a)
+        print(type(a))
+        return 1,2
+
+    @classmethod
+    def build_hpo_from_args(cls, args):
+        """Build a new hpo instance."""
+        return cls(args)
diff --git a/autogl/module/hpo/utils.py b/autogl/module/hpo/utils.py
new file mode 100644
index 0000000..61a4f91
--- /dev/null
+++ b/autogl/module/hpo/utils.py
@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import logging
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import nni.retiarii.nn.pytorch as nn
+from nni.nas.pytorch.mutables import InputChoice, LayerChoice
+
+_logger = logging.getLogger(__name__)
+
+
+def to_device(obj, device):
+    """
+    Move a tensor, tuple, list, or dict onto device.
+    """
+    if torch.is_tensor(obj):
+        return obj.to(device)
+    if isinstance(obj, tuple):
+        return tuple(to_device(t, device) for t in obj)
+    if isinstance(obj, list):
+        return [to_device(t, device) for t in obj]
+    if isinstance(obj, dict):
+        return {k: to_device(v, device) for k, v in obj.items()}
+    if isinstance(obj, (int, float, str)):
+        return obj
+    raise ValueError("'%s' has unsupported type '%s'" % (obj, type(obj)))
+
+
+def to_list(arr):
+    if torch.is_tensor(arr):
+        return arr.cpu().numpy().tolist()
+    if isinstance(arr, np.ndarray):
+        return arr.tolist()
+    if isinstance(arr, (list, tuple)):
+        return list(arr)
+    return arr
+
+
+class AverageMeterGroup:
+    """
+    Average meter group for multiple average meters.
+    """
+
+    def __init__(self):
+        self.meters = OrderedDict()
+
+    def update(self, data):
+        """
+        Update the meter group with a dict of metrics.
+        Non-exist average meters will be automatically created.
+        """
+        for k, v in data.items():
+            if k not in self.meters:
+                self.meters[k] = AverageMeter(k, ":4f")
+            self.meters[k].update(v)
+
+    def __getattr__(self, item):
+        return self.meters[item]
+
+    def __getitem__(self, item):
+        return self.meters[item]
+
+    def __str__(self):
+        return "  ".join(str(v) for v in self.meters.values())
+
+    def summary(self):
+        """
+        Return a summary string of group data.
+        """
+        return "  ".join(v.summary() for v in self.meters.values())
+
+
+class AverageMeter:
+    """
+    Computes and stores the average and current value.
+
+    Parameters
+    ----------
+    name : str
+        Name to display.
+    fmt : str
+        Format string to print the values.
+    """
+
+    def __init__(self, name, fmt=':f'):
+        self.name = name
+        self.fmt = fmt
+        self.reset()
+
+    def reset(self):
+        """
+        Reset the meter.
+        """
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        """
+        Update with value and weight.
+
+        Parameters
+        ----------
+        val : float or int
+            The new value to be accounted in.
+        n : int
+            The weight of the new value.
+        """
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
+        return fmtstr.format(**self.__dict__)
+
+    def summary(self):
+        fmtstr = '{name}: {avg' + self.fmt + '}'
+        return fmtstr.format(**self.__dict__)
+
+
+def _replace_module_with_type(root_module, init_fn, type_name, modules):
+    if modules is None:
+        modules = []
+
+    def apply(m):
+        for name, child in m.named_children():
+            if isinstance(child, type_name):
+                setattr(m, name, init_fn(child))
+                modules.append((child.key, getattr(m, name)))
+            else:
+                apply(child)
+
+    apply(root_module)
+    return modules
+
+
+def replace_layer_choice(root_module, init_fn, modules=None):
+    """
+    Replace layer choice modules with modules that are initiated with init_fn.
+
+    Parameters
+    ----------
+    root_module : nn.Module
+        Root module to traverse.
+    init_fn : Callable
+        Initializing function.
+    modules : dict, optional
+        Update the replaced modules into the dict and check duplicate if provided.
+
+    Returns
+    -------
+    List[Tuple[str, nn.Module]]
+        A list from layer choice keys (names) and replaced modules.
+    """
+    return _replace_module_with_type(root_module, init_fn, (LayerChoice, nn.LayerChoice), modules)
+
+
+def replace_input_choice(root_module, init_fn, modules=None):
+    """
+    Replace input choice modules with modules that are initiated with init_fn.
+
+    Parameters
+    ----------
+    root_module : nn.Module
+        Root module to traverse.
+    init_fn : Callable
+        Initializing function.
+    modules : dict, optional
+        Update the replaced modules into the dict and check duplicate if provided.
+
+    Returns
+    -------
+    List[Tuple[str, nn.Module]]
+        A list from layer choice keys (names) and replaced modules.
+    """
+    return _replace_module_with_type(root_module, init_fn, (InputChoice, nn.InputChoice), modules)
diff --git a/examples/node_classification.py b/examples/node_classification.py
index 939b555..f95c07f 100644
--- a/examples/node_classification.py
+++ b/examples/node_classification.py
@@ -18,7 +18,7 @@ if __name__ == '__main__':
     parser.add_argument('--dataset', default='cora', type=str)
     parser.add_argument('--configs', type=str, default='../configs/nodeclf_gcn_benchmark_small.yml')
     # following arguments will override parameters in the config file
-    parser.add_argument('--hpo', type=str, default='random')
+    parser.add_argument('--hpo', type=str, default='test')
     parser.add_argument('--max_eval', type=int, default=5)
     parser.add_argument('--seed', type=int, default=0)
     parser.add_argument('--device', default=0, type=int)