add bench

test/nas/bench.py

@@ -0,0 +1,392 @@
+"""
+Test file for nas on node classification
+
+AUTOGL_BACKEND=pyg python test/nas/node_classification.py
+AUTOGL_BACKEND=dgl python test/nas/node_classification.py
+
+TODO: make it a unit test file to test all the possible combinations
+"""
+
+import os
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+from autogl.backend import DependentBackend
+
+if DependentBackend.is_dgl():
+    from autogl.module.model.dgl import BaseAutoModel
+    from dgl.data import CoraGraphDataset
+elif DependentBackend.is_pyg():
+    from torch_geometric.datasets import Planetoid
+    from autogl.module.model.pyg import BaseAutoModel
+from autogl.datasets import build_dataset_from_name
+import torch
+import torch.nn.functional as F
+from autogl.module.nas.space.single_path import SinglePathNodeClassificationSpace
+from autogl.module.nas.space.graph_nas import GraphNasNodeClassificationSpace
+from autogl.module.nas.space.graph_nas_macro import GraphNasMacroNodeClassificationSpace
+from autogl.module.nas.estimator.one_shot import OneShotEstimator
+from autogl.module.nas.estimator.train_scratch import TrainEstimator
+from autogl.module.nas.algorithm.agnn_rl import AGNNRL
+from autogl.module.nas.space.autoattend import AutoAttendNodeClassificationSpace
+from autogl.module.nas.backend import bk_feat, bk_label
+from autogl.module.nas.algorithm import Darts, RL, GraphNasRL, Enas, RandomSearch,Spos
+import numpy as np
+from autogl.solver.utils import set_seed
+
+set_seed(202106)
+from autogl.module.nas.space import BaseSpace
+import typing as _typ
+from torch import nn
+class StrModule(nn.Module):
+    def __init__(self, lambd):
+        super().__init__()
+        self.name = lambd
+
+    def forward(self, *args, **kwargs):
+        return self.name
+
+    def __repr__(self):
+        return "{}({})".format(self.__class__.__name__, self.name)
+
+
+gnn_list_proteins = [
+    "gcn",  # GCN
+    "cheb",  # chebnet
+    "arma",
+    "fc",  # skip connection
+    "skip"  # skip connection
+]
+
+gnn_list = [
+    "gat",  # GAT with 2 heads
+    "gcn",  # GCN
+    "gin",  # GIN
+    "cheb",  # chebnet
+    "sage",  # sage
+    "arma",
+    "graph",  
+    "fc",  # skip connection
+    "skip"  # skip connection
+]
+
+
+class Arch:
+    def __init__(self, lk=None, op=None):
+        self.link = lk
+        self.ops = op
+
+    # def random_arch(self):
+    #     self.ops = []
+    #     self.link = random.choice(link_list)
+    #     for i in self.link:
+    #         self.ops.append(random.choice(gnn_list))
+
+    def hash_arch(self, use_proteins = False):
+        lk = self.link
+        op = self.ops
+        if use_proteins:
+            gnn_g = {name: i for i, name in enumerate(gnn_list_proteins)}
+            b = len(gnn_list_proteins) + 1
+        else:
+            gnn_g = {name: i for i, name in enumerate(gnn_list)}
+            b = len(gnn_list) + 1
+        if lk == [0,0,0,0]:
+            lk_hash = 0
+        elif lk == [0,0,0,1]:
+            lk_hash = 1
+        elif lk == [0,0,1,1]:
+            lk_hash = 2
+        elif lk == [0,0,1,2]:
+            lk_hash = 3
+        elif lk == [0,0,1,3]:
+            lk_hash = 4
+        elif lk == [0,1,1,1]:
+            lk_hash = 5
+        elif lk == [0,1,1,2]:
+            lk_hash = 6
+        elif lk == [0,1,2,2]:
+            lk_hash = 7
+        elif lk == [0,1,2,3]:
+            lk_hash = 8
+
+        for i in op:
+            lk_hash = lk_hash * b + gnn_g[i]
+        return lk_hash
+
+    def regularize(self):
+        lk = self.link[:]
+        ops = self.ops[:]
+        if lk == [0,0,0,2]:
+            lk = [0,0,0,1]
+            ops = [ops[1], ops[0], ops[2], ops[3]]
+        elif lk == [0,0,0,3]:
+            lk = [0,0,0,1]
+            ops = [ops[2], ops[0], ops[1], ops[3]]
+        elif lk == [0,0,1,0]:
+            lk = [0,0,0,1]
+            ops = [ops[0], ops[1], ops[3], ops[2]]
+        elif lk == [0,0,2,0]:
+            lk = [0,0,0,1]
+            ops = [ops[1], ops[0], ops[3], ops[2]]
+        elif lk == [0,0,2,1]:
+            lk = [0,0,1,2]
+            ops = [ops[1], ops[0], ops[2], ops[3]]
+        elif lk == [0,0,2,2]:
+            lk = [0,0,1,1]
+            ops = [ops[1], ops[0], ops[2], ops[3]]
+        elif lk == [0,0,2,3]:
+            lk = [0,0,1,3]
+            ops = [ops[1], ops[0], ops[2], ops[3]]
+        elif lk == [0,1,0,0]:
+            lk = [0,0,0,1]
+            ops = [ops[0], ops[2], ops[3], ops[1]]
+        elif lk == [0,1,0,1]:
+            lk = [0,0,1,1]
+            ops = [ops[0], ops[2], ops[1], ops[3]]
+        elif lk == [0,1,0,2]:
+            lk = [0,0,1,3]
+            ops = [ops[0], ops[2], ops[1], ops[3]]
+        elif lk == [0,1,0,3]:
+            lk = [0,0,1,2]
+            ops = [ops[0], ops[2], ops[1], ops[3]]
+        elif lk == [0,1,1,0]:
+            lk = [0,0,1,1]
+            ops = [ops[0], ops[3], ops[1], ops[2]]
+        elif lk == [0,1,1,3]:
+            lk = [0,1,1,2]
+            ops = [ops[0], ops[2], ops[1], ops[3]]
+        elif lk == [0,1,2,0]:
+            lk = [0,0,1,3]
+            ops = [ops[0], ops[3], ops[1], ops[2]]
+        elif lk == [0,1,2,1]:
+            lk = [0,1,1,2]
+            ops = [ops[0], ops[1], ops[3], ops[2]]
+        return Arch(lk, ops)
+
+    def equalpart_sort(self):
+        lk = self.link
+        op = self.ops
+        ops = op[:]
+        def part_sort(ids, ops):
+            gnn_g = {name: i for i, name in enumerate(gnn_list)}
+            opli = [gnn_g[ops[i]] for i in ids]
+            opli.sort()
+            for posid, opid in zip(ids, opli):
+                ops[posid] = gnn_list[opid]
+            return ops
+
+        def sort0012(ops):
+            gnn_g = {name: i for i, name in enumerate(gnn_list)}
+            if gnn_g[op[0]] > gnn_g[op[1]] or op[0] == op[1] and gnn_g[op[2]] > gnn_g[op[3]]:
+                ops = [ops[1], ops[0], ops[3], ops[2]]
+            return ops
+
+        if lk == [0,0,0,0]:
+            ids = [0,1,2,3]
+        elif lk == [0,0,0,1]:
+            ids = [1,2] 
+        elif lk == [0,0,1,1]:
+            ids = [2,3] 
+        elif lk == [0,0,1,2]:
+            ids = None
+            ops = sort0012(ops)
+        elif lk == [0,1,1,1]:
+            ids = [1,2,3] 
+        elif lk == [0,1,2,2]:
+            ids = [2,3] 
+        else:
+            ids = None
+
+        if ids:
+            ops = part_sort(ids, ops)
+
+        self.ops = ops
+
+    def move_skip_op(self):
+        link = self.link[:]
+        ops = self.ops[:]
+        def move_one(k, link, ops):
+            ops = [ops[k]] + ops[:k] + ops[k + 1:]
+            for i, father in enumerate(link):
+                if father == k + 1:
+                    link[i] = link[k]
+                if father <= k:
+                    link[i] = link[i] + 1
+            link = [0] + link[:k] + link[k + 1:]
+            return link, ops
+
+        def check_dim(k, link, ops):
+            # check if a dimension is original dimension
+            while k > -1:
+                if ops[k] != 'skip':
+                    return False
+                k = link[k] - 1
+            return True
+
+        for i in range(len(link)):
+            if ops[i] != 'skip':
+                continue
+            son = False
+            brother = False
+            for j, fa in enumerate(link):
+                if fa == i + 1:
+                    son = True
+                elif j != i and fa == link[i]:
+                    brother = True
+            if son or not brother or check_dim(i, link, ops) and not son:
+                link, ops = move_one(i, link, ops)
+
+        if link == [0,1,2,1]:
+            link = [0,1,1,2]
+            ops = ops[:2] + [ops[3], ops[2]]
+        elif link == [0,1,1,3]:
+            link = [0,1,1,2]
+            ops = [ops[0], ops[2], ops[1], ops[3]]
+
+        #if link not in link_list:
+        #    print(lk, link)
+            
+        self.link = link
+        self.ops = ops
+
+    def valid_hash(self):
+        b = self.regularize()
+        b.move_skip_op()
+        b.equalpart_sort()
+        return b.hash_arch()
+
+    def check_isomorph(self):
+        link, ops = self.link, self.ops
+        linkm = link[:]
+        opsm = ops[:]
+        self.move_skip_op()
+        self.equalpart_sort()
+        #print(self.link, self.ops)
+        return linkm == self.link and opsm == self.ops
+
+import nni
+def map_value(l, label):
+    return nni.retiarii.nn.pytorch.ValueChoice(l, label = label)
+class BenchSpace(BaseSpace):
+    def __init__(
+        self,
+        hidden_dim: _typ.Optional[int] = 64,
+        layer_number: _typ.Optional[int] = 2,
+        dropout: _typ.Optional[float] = 0.9,
+        input_dim: _typ.Optional[int] = None,
+        output_dim: _typ.Optional[int] = None,
+        ops_type = 0
+    ):
+        super().__init__()
+        self.layer_number = layer_number
+        self.hidden_dim = hidden_dim
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.dropout = dropout
+        self.ops_type=ops_type
+
+    def instantiate(
+        self,
+        hidden_dim: _typ.Optional[int] = None,
+        layer_number: _typ.Optional[int] = None,
+        dropout: _typ.Optional[float] = None,
+        input_dim: _typ.Optional[int] = None,
+        output_dim: _typ.Optional[int] = None,
+        ops_type=None
+    ):
+        super().instantiate()
+        self.dropout = dropout or self.dropout
+        self.hidden_dim = hidden_dim or self.hidden_dim
+        self.layer_number = layer_number or self.layer_number
+        self.input_dim = input_dim or self.input_dim
+        self.output_dim = output_dim or self.output_dim
+        self.ops_type = ops_type or self.ops_type
+        self.ops = [gnn_list,gnn_list_proteins][self.ops_type]
+        for layer in range(4):
+            setattr(self,f"in{layer}",self.setInputChoice(layer,n_candidates=layer+1,n_chosen=1,return_mask=False,key=f"in{layer}"))
+            setattr(self,f"op{layer}",self.setLayerChoice(layer,list(map(lambda x:StrModule(x),self.ops)),key=f"op{layer}"))
+        self.dummy=nn.Linear(1,1)
+
+    def forward(self, bench):
+        lks = [getattr(self, "in" + str(i)).selected for i in range(4)]
+        ops = [getattr(self, "op" + str(i)).name for i in range(4)]
+        arch = Arch(lks, ops)
+        h = arch.valid_hash()
+        if h == "88888":
+            return 0
+        return bench[h]['perf']
+
+    def parse_model(self, selection, device) -> BaseAutoModel:
+        return self.wrap().fix(selection)
+
+import os.path as osp
+bench_path='/DATA/DATANAS1/zzy/bench/light'
+import pickle
+
+def light_read(dname):
+    f = open(osp.join(bench_path,f"{dname}.bench"), "rb")
+    bench = pickle.load(f)
+    f.close()
+    return bench
+
+from autogl.module.nas.estimator import BaseEstimator
+from autogl.module.train.evaluation import Acc
+class BenchEstimator(BaseEstimator):
+    def __init__(self, data_name, loss_f="nll_loss", evaluation=[Acc()]):
+        super().__init__(loss_f, evaluation)
+        self.evaluation = evaluation
+        self.bench=light_read(data_name)
+    def infer(self, model: BaseSpace, dataset, mask="train"):
+        perf=model(self.bench)
+        return [perf],0
+
+def run(data_name='cora',algo='graphnas',num_epochs=50,ctrl_steps_aggregate=20):
+    print("Testing backend: {}".format("dgl" if DependentBackend.is_dgl() else "pyg"))
+    if DependentBackend.is_dgl():
+        from autogl.datasets.utils.conversion._to_dgl_dataset import to_dgl_dataset as convert_dataset
+    else:
+        from autogl.datasets.utils.conversion._to_pyg_dataset import to_pyg_dataset as convert_dataset
+
+    # dataset = build_dataset_from_name('cora')
+    # dataset = convert_dataset(dataset)
+    # data = dataset[0]
+
+    # di = bk_feat(data).shape[1]
+    # do = len(np.unique(bk_label(data)))
+    
+    di=2
+    do=2
+    dataset=None
+
+    ops_type=data_name=='proteins'
+
+    space = BenchSpace().cuda()
+    space.instantiate(input_dim=di, output_dim=do,ops_type=ops_type)
+    esti = BenchEstimator(data_name)
+    if algo=='graphnas':
+        algo = GraphNasRL(num_epochs=num_epochs,ctrl_steps_aggregate=ctrl_steps_aggregate)
+    elif algo=='agnn':
+        algo = AGNNRL(guide_type=1,num_epochs=num_epochs,ctrl_steps_aggregate=ctrl_steps_aggregate)
+    else:
+        assert False,f'Not implemented algo {algo}'
+    model = algo.search(space, dataset, esti)
+    result=esti.infer(model._model,None)[0][0]
+    return result
+
+import pandas as pd
+if __name__ == "__main__":
+    data_names='arxiv citeseer computers cora cs photo physics proteins pubmed'.split()
+    algos='graphnas agnn'.split()
+
+    results=[]
+    for data_name in data_names:
+        for algo in algos:
+            print(f'data {data_name} algo {algo}')
+            # metric=run(data_name,algo,2,2)
+            metric=run(data_name,algo,50,20)
+            results.append([algo,data_name,metric])
+
+    df=pd.DataFrame(results,columns='algo data v'.split()).pivot_table(values='v',index='algo',columns='data')
+    print(df.to_string())