Merge branch 'dev' into multiprocess

autogl/datasets/__init__.py

@@ -46,6 +46,7 @@ def register_dataset(name):
 
     return register_dataset_cls
 
+
 from .pyg import (
     AmazonComputersDataset,
     AmazonPhotoDataset,
@@ -96,9 +97,12 @@ from .matlab_matrix import (
     PPIDataset,
 )
 from .modelnet import (
-    ModelNet10, ModelNet40,
-    ModelNet10Train, ModelNet10Test,
-    ModelNet40Train, ModelNet40Test
+    ModelNet10,
+    ModelNet40,
+    ModelNet10Train,
+    ModelNet10Test,
+    ModelNet40Train,
+    ModelNet40Test,
 )
 from .utils import (
     get_label_number,
@@ -110,6 +114,7 @@ from .utils import (
     graph_get_split,
 )
 
+
 def build_dataset(args, path="~/.cache-autogl/"):
     path = osp.join(path, "data", args.dataset)
     path = os.path.expanduser(path)
@@ -120,9 +125,9 @@ def build_dataset_from_name(dataset_name, path="~/.cache-autogl/"):
     path = osp.join(path, "data", dataset_name)
     path = os.path.expanduser(path)
     dataset = DATASET_DICT[dataset_name](path)
-    if 'ogbn' in dataset_name:
-        #dataset.data, dataset.slices = dataset.collate([dataset.data])
-        #dataset.data.num_nodes = dataset.data.num_nodes[0]
+    if "ogbn" in dataset_name:
+        # dataset.data, dataset.slices = dataset.collate([dataset.data])
+        # dataset.data.num_nodes = dataset.data.num_nodes[0]
         if dataset.data.y.shape[-1] == 1:
             dataset.data.y = torch.squeeze(dataset.data.y)
     return dataset

autogl/datasets/modelnet.py

@@ -21,42 +21,42 @@ class ModelNet40(ModelNet):
 @register_dataset("ModelNet10Train")
 class ModelNet10Train(ModelNet):
     def __init__(self, path: str):
-        super(ModelNet10Train, self).__init__(path, '10', train=True)
+        super(ModelNet10Train, self).__init__(path, "10", train=True)
 
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(ModelNet10Train, self).get(idx)
 
 
 @register_dataset("ModelNet10Test")
 class ModelNet10Test(ModelNet):
     def __init__(self, path: str):
-        super(ModelNet10Test, self).__init__(path, '10', train=False)
+        super(ModelNet10Test, self).__init__(path, "10", train=False)
 
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(ModelNet10Test, self).get(idx)
 
 
 @register_dataset("ModelNet40Train")
 class ModelNet40Train(ModelNet):
     def __init__(self, path: str):
-        super(ModelNet40Train, self).__init__(path, '40', train=True)
+        super(ModelNet40Train, self).__init__(path, "40", train=True)
 
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(ModelNet40Train, self).get(idx)
 
 
 @register_dataset("ModelNet40Test")
 class ModelNet40Test(ModelNet):
     def __init__(self, path: str):
-        super(ModelNet40Test, self).__init__(path, '40', train=False)
+        super(ModelNet40Test, self).__init__(path, "40", train=False)
 
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(ModelNet40Test, self).get(idx)

autogl/datasets/ogb.py

@@ -30,15 +30,15 @@ class OGBNproductsDataset(PygNodePropPredDataset):
         split_idx = self.get_idx_split()
         datalist = []
         for d in self:
-            setattr(d, "train_mask", index_to_mask(split_idx['train'], d.y.shape[0]))
-            setattr(d, "val_mask", index_to_mask(split_idx['valid'], d.y.shape[0]))
-            setattr(d, "test_mask", index_to_mask(split_idx['test'], d.y.shape[0]))
+            setattr(d, "train_mask", index_to_mask(split_idx["train"], d.y.shape[0]))
+            setattr(d, "val_mask", index_to_mask(split_idx["valid"], d.y.shape[0]))
+            setattr(d, "test_mask", index_to_mask(split_idx["test"], d.y.shape[0]))
             datalist.append(d)
         self.data, self.slices = self.collate(datalist)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBNproductsDataset, self).get(idx)
 
 
@@ -49,7 +49,9 @@ class OGBNproteinsDataset(PygNodePropPredDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         PygNodePropPredDataset(name=dataset, root=path)
         super(OGBNproteinsDataset, self).__init__(dataset, path)
-        dataset_t = PygNodePropPredDataset(name=dataset, root=path, transform=T.ToSparseTensor())
+        dataset_t = PygNodePropPredDataset(
+            name=dataset, root=path, transform=T.ToSparseTensor()
+        )
 
         # Move edge features to node features.
         self.data.x = dataset_t[0].adj_t.mean(dim=1)
@@ -61,15 +63,15 @@ class OGBNproteinsDataset(PygNodePropPredDataset):
         split_idx = self.get_idx_split()
         datalist = []
         for d in self:
-            setattr(d, "train_mask", index_to_mask(split_idx['train'], d.y.shape[0]))
-            setattr(d, "val_mask", index_to_mask(split_idx['valid'], d.y.shape[0]))
-            setattr(d, "test_mask", index_to_mask(split_idx['test'], d.y.shape[0]))
+            setattr(d, "train_mask", index_to_mask(split_idx["train"], d.y.shape[0]))
+            setattr(d, "val_mask", index_to_mask(split_idx["valid"], d.y.shape[0]))
+            setattr(d, "test_mask", index_to_mask(split_idx["test"], d.y.shape[0]))
             datalist.append(d)
         self.data, self.slices = self.collate(datalist)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBNproteinsDataset, self).get(idx)
 
 
@@ -86,15 +88,15 @@ class OGBNarxivDataset(PygNodePropPredDataset):
 
         datalist = []
         for d in self:
-            setattr(d, "train_mask", index_to_mask(split_idx['train'], d.y.shape[0]))
-            setattr(d, "val_mask", index_to_mask(split_idx['valid'], d.y.shape[0]))
-            setattr(d, "test_mask", index_to_mask(split_idx['test'], d.y.shape[0]))
+            setattr(d, "train_mask", index_to_mask(split_idx["train"], d.y.shape[0]))
+            setattr(d, "val_mask", index_to_mask(split_idx["valid"], d.y.shape[0]))
+            setattr(d, "test_mask", index_to_mask(split_idx["test"], d.y.shape[0]))
             datalist.append(d)
         self.data, self.slices = self.collate(datalist)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBNarxivDataset, self).get(idx)
 
 
@@ -110,15 +112,15 @@ class OGBNpapers100MDataset(PygNodePropPredDataset):
         split_idx = self.get_idx_split()
         datalist = []
         for d in self:
-            setattr(d, "train_mask", index_to_mask(split_idx['train'], d.y.shape[0]))
-            setattr(d, "val_mask", index_to_mask(split_idx['valid'], d.y.shape[0]))
-            setattr(d, "test_mask", index_to_mask(split_idx['test'], d.y.shape[0]))
+            setattr(d, "train_mask", index_to_mask(split_idx["train"], d.y.shape[0]))
+            setattr(d, "val_mask", index_to_mask(split_idx["valid"], d.y.shape[0]))
+            setattr(d, "test_mask", index_to_mask(split_idx["test"], d.y.shape[0]))
             datalist.append(d)
         self.data, self.slices = self.collate(datalist)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBNpapers100MDataset, self).get(idx)
 
 
@@ -134,9 +136,10 @@ class OGBNmagDataset(PygNodePropPredDataset):
         rel_data = self[0]
         # We are only interested in paper <-> paper relations.
         self.data = Data(
-            x=rel_data.x_dict['paper'],
-            edge_index=rel_data.edge_index_dict[('paper', 'cites', 'paper')],
-            y=rel_data.y_dict['paper'])
+            x=rel_data.x_dict["paper"],
+            edge_index=rel_data.edge_index_dict[("paper", "cites", "paper")],
+            y=rel_data.y_dict["paper"],
+        )
 
         # self.data = T.ToSparseTensor()(data)
         # self[0].adj_t = self[0].adj_t.to_symmetric()
@@ -147,15 +150,15 @@ class OGBNmagDataset(PygNodePropPredDataset):
 
         datalist = []
         for d in self:
-            setattr(d, "train_mask", index_to_mask(split_idx['train'], d.y.shape[0]))
-            setattr(d, "val_mask", index_to_mask(split_idx['valid'], d.y.shape[0]))
-            setattr(d, "test_mask", index_to_mask(split_idx['test'], d.y.shape[0]))
+            setattr(d, "train_mask", index_to_mask(split_idx["train"], d.y.shape[0]))
+            setattr(d, "val_mask", index_to_mask(split_idx["valid"], d.y.shape[0]))
+            setattr(d, "test_mask", index_to_mask(split_idx["test"], d.y.shape[0]))
             datalist.append(d)
         self.data, self.slices = self.collate(datalist)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBNmagDataset, self).get(idx)
 
 
@@ -171,10 +174,10 @@ class OGBGmolhivDataset(PygGraphPropPredDataset):
         super(OGBGmolhivDataset, self).__init__(dataset, path)
         setattr(OGBGmolhivDataset, "metric", "ROC-AUC")
         setattr(OGBGmolhivDataset, "loss", "binary_cross_entropy_with_logits")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBGmolhivDataset, self).get(idx)
 
 
@@ -187,10 +190,10 @@ class OGBGmolpcbaDataset(PygGraphPropPredDataset):
         super(OGBGmolpcbaDataset, self).__init__(dataset, path)
         setattr(OGBGmolpcbaDataset, "metric", "AP")
         setattr(OGBGmolpcbaDataset, "loss", "binary_cross_entropy_with_logits")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBGmolpcbaDataset, self).get(idx)
 
 
@@ -203,10 +206,10 @@ class OGBGppaDataset(PygGraphPropPredDataset):
         super(OGBGppaDataset, self).__init__(dataset, path)
         setattr(OGBGppaDataset, "metric", "Accuracy")
         setattr(OGBGppaDataset, "loss", "cross_entropy")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBGppaDataset, self).get(idx)
 
 
@@ -219,10 +222,10 @@ class OGBGcodeDataset(PygGraphPropPredDataset):
         super(OGBGcodeDataset, self).__init__(dataset, path)
         setattr(OGBGcodeDataset, "metric", "F1 score")
         setattr(OGBGcodeDataset, "loss", "cross_entropy")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBGcodeDataset, self).get(idx)
 
 
@@ -238,10 +241,10 @@ class OGBLppaDataset(PygLinkPropPredDataset):
         super(OGBLppaDataset, self).__init__(dataset, path)
         setattr(OGBLppaDataset, "metric", "Hits@100")
         setattr(OGBLppaDataset, "loss", "pos_neg_loss")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBLppaDataset, self).get(idx)
 
 
@@ -254,10 +257,10 @@ class OGBLcollabDataset(PygLinkPropPredDataset):
         super(OGBLcollabDataset, self).__init__(dataset, path)
         setattr(OGBLcollabDataset, "metric", "Hits@50")
         setattr(OGBLcollabDataset, "loss", "pos_neg_loss")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBLcollabDataset, self).get(idx)
 
 
@@ -270,10 +273,10 @@ class OGBLddiDataset(PygLinkPropPredDataset):
         super(OGBLddiDataset, self).__init__(dataset, path)
         setattr(OGBLddiDataset, "metric", "Hits@20")
         setattr(OGBLddiDataset, "loss", "pos_neg_loss")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBLddiDataset, self).get(idx)
 
 
@@ -286,10 +289,10 @@ class OGBLcitationDataset(PygLinkPropPredDataset):
         super(OGBLcitationDataset, self).__init__(dataset, path)
         setattr(OGBLcitationDataset, "metric", "MRR")
         setattr(OGBLcitationDataset, "loss", "pos_neg_loss")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBLcitationDataset, self).get(idx)
 
 
@@ -302,10 +305,10 @@ class OGBLwikikgDataset(PygLinkPropPredDataset):
         super(OGBLwikikgDataset, self).__init__(dataset, path)
         setattr(OGBLwikikgDataset, "metric", "MRR")
         setattr(OGBLwikikgDataset, "loss", "pos_neg_loss")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBLwikikgDataset, self).get(idx)
 
 
@@ -318,8 +321,8 @@ class OGBLbiokgDataset(PygLinkPropPredDataset):
         super(OGBLbiokgDataset, self).__init__(dataset, path)
         setattr(OGBLbiokgDataset, "metric", "MRR")
         setattr(OGBLbiokgDataset, "loss", "pos_neg_loss")
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(OGBLbiokgDataset, self).get(idx)

autogl/datasets/pyg.py

@@ -1,6 +1,7 @@
 import os.path as osp
 
 import torch
+
 # import torch_geometric.transforms as T
 from torch_geometric.datasets import (
     Planetoid,
@@ -21,10 +22,10 @@ class AmazonComputersDataset(Amazon):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Amazon(path, dataset)
         super(AmazonComputersDataset, self).__init__(path, dataset)
-        
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(AmazonComputersDataset, self).get(idx)
 
 
@@ -35,10 +36,10 @@ class AmazonPhotoDataset(Amazon):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Amazon(path, dataset)
         super(AmazonPhotoDataset, self).__init__(path, dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(AmazonPhotoDataset, self).get(idx)
 
 
@@ -49,10 +50,10 @@ class CoauthorPhysicsDataset(Coauthor):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Coauthor(path, dataset)
         super(CoauthorPhysicsDataset, self).__init__(path, dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(CoauthorPhysicsDataset, self).get(idx)
 
 
@@ -63,10 +64,10 @@ class CoauthorCSDataset(Coauthor):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Coauthor(path, dataset)
         super(CoauthorCSDataset, self).__init__(path, dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(CoauthorCSDataset, self).get(idx)
 
 
@@ -77,10 +78,10 @@ class CoraDataset(Planetoid):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Planetoid(path, dataset)
         super(CoraDataset, self).__init__(path, dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(CoraDataset, self).get(idx)
 
 
@@ -91,10 +92,10 @@ class CiteSeerDataset(Planetoid):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Planetoid(path, dataset)
         super(CiteSeerDataset, self).__init__(path, dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(CiteSeerDataset, self).get(idx)
 
 
@@ -105,10 +106,10 @@ class PubMedDataset(Planetoid):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Planetoid(path, dataset)
         super(PubMedDataset, self).__init__(path, dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(PubMedDataset, self).get(idx)
 
 
@@ -119,10 +120,10 @@ class RedditDataset(Reddit):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         Reddit(path)
         super(RedditDataset, self).__init__(path)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(RedditDataset, self).get(idx)
 
 
@@ -135,8 +136,8 @@ class MUTAGDataset(TUDataset):
         super(MUTAGDataset, self).__init__(path, name=dataset)
 
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(MUTAGDataset, self).get(idx)
 
 
@@ -147,10 +148,10 @@ class IMDBBinaryDataset(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(IMDBBinaryDataset, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(IMDBBinaryDataset, self).get(idx)
 
 
@@ -161,10 +162,10 @@ class IMDBMultiDataset(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(IMDBMultiDataset, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(IMDBMultiDataset, self).get(idx)
 
 
@@ -175,10 +176,10 @@ class CollabDataset(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(CollabDataset, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(CollabDataset, self).get(idx)
 
 
@@ -189,10 +190,10 @@ class ProteinsDataset(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(ProteinsDataset, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(ProteinsDataset, self).get(idx)
 
 
@@ -203,10 +204,10 @@ class REDDITBinary(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(REDDITBinary, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(REDDITBinary, self).get(idx)
 
 
@@ -217,10 +218,10 @@ class REDDITMulti5K(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(REDDITMulti5K, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(REDDITMulti5K, self).get(idx)
 
 
@@ -231,10 +232,10 @@ class REDDITMulti12K(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(REDDITMulti12K, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(REDDITMulti12K, self).get(idx)
 
 
@@ -247,8 +248,8 @@ class PTCMRDataset(TUDataset):
         super(PTCMRDataset, self).__init__(path, name=dataset)
 
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(PTCMRDataset, self).get(idx)
 
 
@@ -259,10 +260,10 @@ class NCI1Dataset(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(NCI1Dataset, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(NCI1Dataset, self).get(idx)
 
 
@@ -273,10 +274,10 @@ class NCI109Dataset(TUDataset):
         # path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset)
         TUDataset(path, name=dataset)
         super(NCI109Dataset, self).__init__(path, name=dataset)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(NCI109Dataset, self).get(idx)
 
 
@@ -298,10 +299,10 @@ class ENZYMES(TUDataset):
             return data
         else:
             return self.index_select(idx)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(ENZYMES, self).get(idx)
 
 
@@ -342,8 +343,8 @@ class QM9Dataset(QM9):
         if not osp.exists(path):
             QM9(path)
         super(QM9Dataset, self).__init__(path)
-    
+
     def get(self, idx):
-        if hasattr(self, '__data_list__'):
-            delattr(self, '__data_list__')
+        if hasattr(self, "__data_list__"):
+            delattr(self, "__data_list__")
         return super(QM9Dataset, self).get(idx)

autogl/datasets/utils.py

@@ -38,7 +38,7 @@ def random_splits_mask(dataset, train_ratio=0.2, val_ratio=0.4, seed=None):
     assert (
         train_ratio + val_ratio <= 1
     ), "the sum of train_ratio and val_ratio is larger than 1"
-    _dataset=[d for d in dataset]
+    _dataset = [d for d in dataset]
     for data in _dataset:
         r_s = torch.get_rng_state()
         if torch.cuda.is_available():
@@ -65,8 +65,8 @@ def random_splits_mask(dataset, train_ratio=0.2, val_ratio=0.4, seed=None):
             torch.cuda.set_rng_state(r_s_cuda)
 
     dataset.data, dataset.slices = dataset.collate(_dataset)
-    if hasattr(dataset, '__data_list__'):
-        delattr(dataset, '__data_list__')
+    if hasattr(dataset, "__data_list__"):
+        delattr(dataset, "__data_list__")
     # while type(dataset.data.num_nodes) == list:
     #    dataset.data.num_nodes = dataset.data.num_nodes[0]
     # dataset.data.num_nodes = dataset.data.num_nodes[0]
@@ -171,8 +171,8 @@ def random_splits_mask_class(
         setattr(d, "test_mask", data.test_mask)
         datalist.append(d)
     dataset.data, dataset.slices = dataset.collate(datalist)
-    if hasattr(dataset, '__data_list__'):
-        delattr(dataset, '__data_list__')
+    if hasattr(dataset, "__data_list__"):
+        delattr(dataset, "__data_list__")
     # while type(dataset.data.num_nodes) == list:
     #     dataset.data.num_nodes = dataset.data.num_nodes[0]
     # dataset.data.num_nodes = dataset.data.num_nodes[0]

autogl/module/ensemble/__init__.py

@@ -16,9 +16,11 @@ def register_ensembler(name):
 
     return register_ensembler_cls
 
+
 from .voting import Voting
 from .stacking import Stacking
 
+
 def build_ensembler_from_name(name: str) -> BaseEnsembler:
     """
     Parameters

autogl/module/ensemble/voting.py

@@ -85,7 +85,7 @@ class Voting(BaseEnsembler):
         weights = weights / np.sum(weights)
 
         return np.average(predictions, axis=0, weights=weights)
-    
+
     def _specify_weights(self, predictions, label, feval):
         ensemble_prediction = []
         combinations = []

autogl/module/feature/__init__.py

@@ -24,6 +24,7 @@ def register_feature(name):
 
     return register_feature_cls
 
+
 from .auto_feature import AutoFeatureEngineer
 from .base import BaseFeatureEngineer
 

autogl/module/feature/auto_feature.py

@@ -12,6 +12,7 @@ from . import register_feature
 
 from ...utils import get_logger
 import torch
+
 LOGGER = get_logger("Feature")
 
 
@@ -28,13 +29,15 @@ class Onlyconst(BaseFeatureEngineer):
     r"""it is a dummy feature engineer , which directly returns identical data"""
 
     def __init__(self, *args, **kwargs):
-        super(Onlyconst, self).__init__(data_t='tensor',multigraph=True, *args, **kwargs)
+        super(Onlyconst, self).__init__(
+            data_t="tensor", multigraph=True, *args, **kwargs
+        )
 
     def _transform(self, data):
-        if 'x' in data:
+        if "x" in data:
             data.x = torch.ones((data.x.shape[0], 1))
         else:
-            data.x= torch.ones((torch.unique(data.edge_index).shape[0],1))
+            data.x = torch.ones((torch.unique(data.edge_index).shape[0], 1))
         return data
 
 

autogl/module/feature/base.py

@@ -68,12 +68,13 @@ class BaseFeatureAtom:
         elif self._data_t == "nx":
             if not hasattr(data, "G") or data.G is None:
                 data.G = to_networkx(data, to_undirected=True)
-    def _adjust_to_tensor(self,data):
+
+    def _adjust_to_tensor(self, data):
         if self._data_t == "tensor":
             pass
         else:
             data_np2tensor(data)
-        
+
     def _preprocess(self, data):
         pass
 
@@ -114,7 +115,6 @@ class BaseFeatureAtom:
                     p._adjust_to_tensor(datai)
                     _dataset[i] = datai
                 dataset = self._rebuild(dataset, _dataset)
-                
 
     def transform(self, dataset, inplace=True):
         r"""transform dataset inplace or not w.r.t bool argument ``inplace``"""
@@ -131,7 +131,7 @@ class BaseFeatureAtom:
                     datai = p._transform(datai)
                     p._postprocess(datai)
                     p._adjust_to_tensor(datai)
-                    _dataset[i] = datai    
+                    _dataset[i] = datai
                 dataset = self._rebuild(dataset, _dataset)
         dataset.data = data_np2tensor(dataset.data)
         return dataset

autogl/module/feature/generators/base.py

@@ -4,8 +4,10 @@ from ..base import BaseFeatureAtom
 
 
 class BaseGenerator(BaseFeatureAtom):
-    def __init__(self, data_t="np", multigraph=True,**kwargs):
-        super(BaseGenerator, self).__init__(data_t=data_t, multigraph=multigraph,**kwargs)
+    def __init__(self, data_t="np", multigraph=True, **kwargs):
+        super(BaseGenerator, self).__init__(
+            data_t=data_t, multigraph=multigraph, **kwargs
+        )
 
 
 @register_feature("onehot")

autogl/module/feature/generators/pyg.py

@@ -78,12 +78,14 @@ class PYGOneHotDegree(PYGGenerator):
     def __init__(self, max_degree=1000):
         super(PYGOneHotDegree, self).__init__(max_degree=max_degree)
 
+    """
     def _transform(self, data):
-        idx, x = data.edge_index[0], data.x
-        deg = degree(idx, data.num_nodes, dtype=torch.long)
-        self._kwargs["max_degree"] = np.min(
-            [self._kwargs["max_degree"], torch.max(deg).numpy()]
-        )
+        #idx, x = data.edge_index[0], data.x
+        #deg = degree(idx, data.num_nodes, dtype=torch.long)
+        #self._kwargs["max_degree"] = np.min(
+        #    [self._kwargs["max_degree"], torch.max(deg).numpy()]
+        #)
         dsc = self.extract(data)
         data.x = torch.cat([data.x, dsc], dim=1)
         return data
+    """

autogl/module/feature/selectors/base.py

@@ -3,8 +3,10 @@ import numpy as np
 
 
 class BaseSelector(BaseFeatureAtom):
-    def __init__(self, data_t="np", multigraph=False,**kwargs):
-        super(BaseSelector, self).__init__(data_t=data_t, multigraph=multigraph,**kwargs)
+    def __init__(self, data_t="np", multigraph=False, **kwargs):
+        super(BaseSelector, self).__init__(
+            data_t=data_t, multigraph=multigraph, **kwargs
+        )
         self._sel = None
 
     def _transform(self, data):

autogl/module/feature/subgraph/base.py

@@ -3,11 +3,12 @@ import numpy as np
 import torch
 from .. import register_feature
 
-@register_feature('subgraph')
+
+@register_feature("subgraph")
 class BaseSubgraph(BaseFeatureAtom):
-    def __init__(self, data_t="np", multigraph=True,**kwargs):
+    def __init__(self, data_t="np", multigraph=True, **kwargs):
         super(BaseSubgraph, self).__init__(
-            data_t=data_t, multigraph=multigraph, subgraph=True,**kwargs
+            data_t=data_t, multigraph=multigraph, subgraph=True, **kwargs
         )
 
     def _preprocess(self, data):
@@ -16,5 +17,3 @@ class BaseSubgraph(BaseFeatureAtom):
 
     def _postprocess(self, data):
         pass
-
-

autogl/module/hpo/advisorbase.py

@@ -43,7 +43,9 @@ class AdvisorBaseHPOptimizer(BaseHPOptimizer):
         self.xs = []
         self.best_id = None
         self.best_trainer = None
-        space = trainer.hyper_parameter_space
+        space = (
+            trainer.hyper_parameter_space + trainer.get_model().hyper_parameter_space
+        )
         current_config = self._encode_para(space)
 
         for i in range(slaves):
@@ -129,7 +131,9 @@ class AdvisorBaseHPOptimizer(BaseHPOptimizer):
 
         self.feval_name = trainer.get_feval(return_major=True).get_eval_name()
         self.is_higher_better = trainer.get_feval(return_major=True).is_higher_better()
-        space = trainer.hyper_parameter_space
+        space = (
+            trainer.hyper_parameter_space + trainer.get_model().hyper_parameter_space
+        )
         current_space = self._encode_para(space)
         self._setUp(current_space)
 

autogl/module/hpo/autone.py

@@ -17,13 +17,15 @@ from torch_geometric.data import GraphSAINTRandomWalkSampler
 from ..feature.subgraph.nx import NxSubgraph, NxLargeCliqueSize
 from ..feature.subgraph import nx, SgNetLSD
 
-from torch_geometric.data import InMemoryDataset 
+from torch_geometric.data import InMemoryDataset
+
 
 class _MyDataset(InMemoryDataset):
     def __init__(self, datalist) -> None:
         super().__init__()
         self.data, self.slices = self.collate(datalist)
 
+
 @register_hpo("autone")
 class AutoNE(BaseHPOptimizer):
     """
@@ -59,7 +61,9 @@ class AutoNE(BaseHPOptimizer):
         """
         self.feval_name = trainer.get_feval(return_major=True).get_eval_name()
         self.is_higher_better = trainer.get_feval(return_major=True).is_higher_better()
-        space = trainer.hyper_parameter_space + trainer.model.hyper_parameter_space
+        space = (
+            trainer.hyper_parameter_space + trainer.get_model().hyper_parameter_space
+        )
         current_space = self._encode_para(space)
 
         def sample_subgraph(whole_data):
@@ -73,17 +77,17 @@ class AutoNE(BaseHPOptimizer):
             )
             results = []
             for data in loader:
-                in_dataset= _MyDataset([data])
+                in_dataset = _MyDataset([data])
                 results.append(in_dataset)
             return results
 
         func = SgNetLSD()
 
         def get_wne(graph):
-            graph=func.fit_transform(graph)
+            graph = func.fit_transform(graph)
             # transform = nx.NxSubgraph.compose(map(lambda x: x(), nx.NX_EXTRACTORS))
             # print(type(graph))
-            #gf = transform.fit_transform(graph).data.gf
+            # gf = transform.fit_transform(graph).data.gf
             gf = graph.data.gf
             fin = list(gf[0]) + list(map(lambda x: float(x), gf[1:]))
             return fin

autogl/module/hpo/suggestion/__init__.py

@@ -1 +1 @@
-# Files in this folder are reproduced from https://github.com/tobegit3hub/advisor with some changes.
+# Files in this folder are reproduced from https://github.com/tobegit3hub/advisor with some changes.

autogl/module/model/__init__.py

@@ -1,35 +1,7 @@
-import importlib
-import os
-
-MODEL_DICT = {}
-
-
-def register_model(name):
-    def register_model_cls(cls):
-        if name in MODEL_DICT:
-            raise ValueError("Cannot register duplicate trainer ({})".format(name))
-        if not issubclass(cls, BaseModel):
-            raise ValueError(
-                "Trainer ({}: {}) must extend BaseModel".format(name, cls.__name__)
-            )
-        MODEL_DICT[name] = cls
-        return cls
-
-    return register_model_cls
-
+from ._model_registry import MODEL_DICT, ModelUniversalRegistry, register_model
 from .base import BaseModel
 from .topkpool import AutoTopkpool
-from .graphsage import AutoSAGE
+from .graph_sage import AutoSAGE
 from .gcn import AutoGCN
 from .gat import AutoGAT
 from .gin import AutoGIN
-
-
-__all__ = [
-    "BaseModel",
-    "AutoTopkpool",
-    "AutoSAGE",
-    "AutoGCN",
-    "AutoGAT",
-    "AutoGIN",
-]

autogl/module/model/_model_registry.py

@@ -0,0 +1,28 @@
+import typing as _typing
+from .base import BaseModel
+
+MODEL_DICT: _typing.Dict[str, _typing.Type[BaseModel]] = {}
+
+
+def register_model(name):
+    def register_model_cls(cls):
+        if name in MODEL_DICT:
+            raise ValueError("Cannot register duplicate trainer ({})".format(name))
+        if not issubclass(cls, BaseModel):
+            raise ValueError(
+                "Trainer ({}: {}) must extend BaseModel".format(name, cls.__name__)
+            )
+        MODEL_DICT[name] = cls
+        return cls
+    
+    return register_model_cls
+
+
+class ModelUniversalRegistry:
+    @classmethod
+    def get_model(cls, name: str) -> _typing.Type[BaseModel]:
+        if type(name) != str:
+            raise TypeError
+        if name not in MODEL_DICT:
+            raise KeyError
+        return MODEL_DICT.get(name)

autogl/module/model/base.py

@@ -43,6 +43,11 @@ class BaseModel(torch.nn.Module):
     def forward(self):
         pass
 
+    def to(self, device):
+        if isinstance(device, (str, torch.device)):
+            self.device = device
+        return super().to(device)
+
     def from_hyper_parameter(self, hp):
         ret_self = self.__class__(
             num_features=self.num_features,

autogl/module/model/gat.py

@@ -21,9 +21,22 @@ class GAT(torch.nn.Module):
         self.args = args
         self.num_layer = int(self.args["num_layers"])
 
-        missing_keys = list(set(["features_num", "num_class", "num_layers", "hidden", "heads", "dropout", "act"]) - set(self.args.keys()))
+        missing_keys = list(
+            set(
+                [
+                    "features_num",
+                    "num_class",
+                    "num_layers",
+                    "hidden",
+                    "heads",
+                    "dropout",
+                    "act",
+                ]
+            )
+            - set(self.args.keys())
+        )
         if len(missing_keys) > 0:
-            raise Exception("Missing keys: %s." % ','.join(missing_keys))
+            raise Exception("Missing keys: %s." % ",".join(missing_keys))
 
         if not self.num_layer == len(self.args["hidden"]) + 1:
             LOGGER.warn("Warning: layer size does not match the length of hidden units")

autogl/module/model/gcn.py

@@ -21,9 +21,12 @@ class GCN(torch.nn.Module):
         self.args = args
         self.num_layer = int(self.args["num_layers"])
 
-        missing_keys = list(set(["features_num", "num_class", "num_layers", "hidden", "dropout", "act"]) - set(self.args.keys()))
+        missing_keys = list(
+            set(["features_num", "num_class", "num_layers", "hidden", "dropout", "act"])
+            - set(self.args.keys())
+        )
         if len(missing_keys) > 0:
-            raise Exception("Missing keys: %s." % ','.join(missing_keys))
+            raise Exception("Missing keys: %s." % ",".join(missing_keys))
 
         if not self.num_layer == len(self.args["hidden"]) + 1:
             LOGGER.warn("Warning: layer size does not match the length of hidden units")

autogl/module/model/gin.py

@@ -25,14 +25,27 @@ class GIN(torch.nn.Module):
         self.num_layer = int(self.args["num_layers"])
         assert self.num_layer > 2, "Number of layers in GIN should not less than 3"
 
-        missing_keys = list(set(["features_num", "num_class", "num_graph_features",
-                    "num_layers", "hidden", "dropout", "act",
-                    "mlp_layers", "eps"]) - set(self.args.keys()))
+        missing_keys = list(
+            set(
+                [
+                    "features_num",
+                    "num_class",
+                    "num_graph_features",
+                    "num_layers",
+                    "hidden",
+                    "dropout",
+                    "act",
+                    "mlp_layers",
+                    "eps",
+                ]
+            )
+            - set(self.args.keys())
+        )
         if len(missing_keys) > 0:
-            raise Exception("Missing keys: %s." % ','.join(missing_keys))
-        if not self.num_layer == len(self.args['hidden']) + 1:
-            LOGGER.warn('Warning: layer size does not match the length of hidden units')
-        self.num_graph_features = self.args['num_graph_features']
+            raise Exception("Missing keys: %s." % ",".join(missing_keys))
+        if not self.num_layer == len(self.args["hidden"]) + 1:
+            LOGGER.warn("Warning: layer size does not match the length of hidden units")
+        self.num_graph_features = self.args["num_graph_features"]
 
         if self.args["act"] == "leaky_relu":
             act = LeakyReLU()

autogl/module/model/graph_sage.py

@@ -0,0 +1,124 @@
+import typing as _typing
+import torch
+import torch.nn.functional as F
+from torch_geometric.nn.conv import SAGEConv
+
+from . import register_model
+from .base import BaseModel, activate_func
+
+
+class GraphSAGE(torch.nn.Module):
+    def __init__(
+            self, num_features: int, num_classes: int,
+            hidden_features: _typing.Sequence[int],
+            dropout: float, activation_name: str,
+            aggr: str = "mean", **kwargs
+    ):
+        super(GraphSAGE, self).__init__()
+        if type(aggr) != str:
+            raise TypeError
+        if aggr not in ("add", "max", "mean"):
+            aggr = "mean"
+        
+        self.__convolution_layers: torch.nn.ModuleList = torch.nn.ModuleList()
+        
+        num_layers: int = len(hidden_features) + 1
+        if num_layers == 1:
+            self.__convolution_layers.append(
+                SAGEConv(num_features, num_classes, aggr=aggr)
+            )
+        else:
+            self.__convolution_layers.append(
+                SAGEConv(num_features, hidden_features[0], aggr=aggr)
+            )
+            for i in range(len(hidden_features)):
+                if i + 1 < len(hidden_features):
+                    self.__convolution_layers.append(
+                        SAGEConv(hidden_features[i], hidden_features[i + 1], aggr=aggr)
+                    )
+                else:
+                    self.__convolution_layers.append(
+                        SAGEConv(hidden_features[i], num_classes, aggr=aggr)
+                    )
+        self.__dropout: float = dropout
+        self.__activation_name: str = activation_name
+    
+    def __full_forward(self, data):
+        x: torch.Tensor = getattr(data, "x")
+        edge_index: torch.Tensor = getattr(data, "edge_index")
+        for layer_index in range(len(self.__convolution_layers)):
+            x: torch.Tensor = self.__convolution_layers[layer_index](x, edge_index)
+            if layer_index + 1 < len(self.__convolution_layers):
+                x = activate_func(x, self.__activation_name)
+                x = F.dropout(x, p=self.__dropout, training=self.training)
+        return F.log_softmax(x, dim=1)
+    
+    def __distributed_forward(self, data):
+        x: torch.Tensor = getattr(data, "x")
+        edge_indexes: _typing.Sequence[torch.Tensor] = getattr(data, "edge_indexes")
+        if len(edge_indexes) != len(self.__convolution_layers):
+            raise AttributeError
+        for layer_index in range(len(self.__convolution_layers)):
+            x: torch.Tensor = self.__convolution_layers[layer_index](x, edge_indexes[layer_index])
+            if layer_index + 1 < len(self.__convolution_layers):
+                x = activate_func(x, self.__activation_name)
+                x = F.dropout(x, p=self.__dropout, training=self.training)
+        return F.log_softmax(x, dim=1)
+    
+    def forward(self, data):
+        if (
+                hasattr(data, "edge_indexes") and
+                isinstance(getattr(data, "edge_indexes"), _typing.Sequence) and
+                len(getattr(data, "edge_indexes")) == len(self.__convolution_layers)
+        ):
+            return self.__distributed_forward(data)
+        else:
+            return self.__full_forward(data)
+
+
+@register_model("sage")
+class AutoSAGE(BaseModel):
+    def __init__(
+            self, num_features: int = 1, num_classes: int = 1,
+            device: _typing.Optional[torch.device] = torch.device("cpu"),
+            init: bool = False, **kwargs
+    ):
+        super(AutoSAGE, self).__init__(init)
+        self.__num_features: int = num_features
+        self.__num_classes: int = num_classes
+        self.__device: torch.device = device if device is not None else torch.device("cpu")
+        
+        self.hyperparams = {
+            "num_layers": 3,
+            "hidden": [64, 32],
+            "dropout": 0.5,
+            "act": "relu",
+            "aggr": "mean",
+        }
+        self.params = {
+            "num_features": self.__num_features,
+            "num_classes": self.__num_classes
+        }
+        
+        self._model: GraphSAGE = GraphSAGE(
+            self.__num_features, self.__num_classes, [64, 32], 0.5, "relu"
+        )
+        
+        self._initialized: bool = False
+        if init:
+            self.initialize()
+    
+    @property
+    def model(self) -> GraphSAGE:
+        return self._model
+    
+    def initialize(self):
+        """ Initialize model """
+        if not self._initialized:
+            self._model: GraphSAGE = GraphSAGE(
+                self.__num_features, self.__num_classes,
+                hidden_features=self.hyperparams["hidden"],
+                activation_name=self.hyperparams["act"],
+                **self.hyperparams
+            )
+            self._initialized = True

autogl/module/model/graphsage.py

@@ -113,11 +113,23 @@ class GraphSAGE(torch.nn.Module):
         if not self.num_layer == len(self.args["hidden"]) + 1:
             LOGGER.warn("Warning: layer size does not match the length of hidden units")
 
-        missing_keys = list(set(["features_num", "num_class", "num_layers",
-                    "hidden", "dropout", "act", "agg"]) - set(self.args.keys()))
+        missing_keys = list(
+            set(
+                [
+                    "features_num",
+                    "num_class",
+                    "num_layers",
+                    "hidden",
+                    "dropout",
+                    "act",
+                    "agg",
+                ]
+            )
+            - set(self.args.keys())
+        )
         if len(missing_keys) > 0:
-            raise Exception("Missing keys: %s." % ','.join(missing_keys))
-        
+            raise Exception("Missing keys: %s." % ",".join(missing_keys))
+
         self.convs = torch.nn.ModuleList()
         self.convs.append(
             SAGEConv(self.args["features_num"], self.args["hidden"][0], aggr=agg)
@@ -160,7 +172,7 @@ class GraphSAGE(torch.nn.Module):
         return F.log_softmax(x, dim=1)
 
 
-@register_model("sage")
+# @register_model("sage")
 class AutoSAGE(BaseModel):
     r"""
     AutoSAGE. The model used in this automodel is GraphSAGE, i.e., the GraphSAGE from the `"Inductive Representation Learning on

autogl/module/model/topkpool.py

@@ -21,10 +21,21 @@ class Topkpool(torch.nn.Module):
         super(Topkpool, self).__init__()
         self.args = args
 
-        missing_keys = list(set(["features_num", "num_class", "num_graph_features",
-                    "ratio", "dropout", "act"]) - set(self.args.keys()))
+        missing_keys = list(
+            set(
+                [
+                    "features_num",
+                    "num_class",
+                    "num_graph_features",
+                    "ratio",
+                    "dropout",
+                    "act",
+                ]
+            )
+            - set(self.args.keys())
+        )
         if len(missing_keys) > 0:
-            raise Exception("Missing keys: %s." % ','.join(missing_keys))
+            raise Exception("Missing keys: %s." % ",".join(missing_keys))
 
         self.num_features = self.args["features_num"]
         self.num_classes = self.args["num_class"]

autogl/module/train/__init__.py

@@ -1,8 +1,14 @@
 import importlib
 import os
-from .base import BaseTrainer, Evaluation, EarlyStopping
 
 TRAINER_DICT = {}
+EVALUATE_DICT = {}
+from .base import (
+    BaseTrainer,
+    Evaluation,
+    BaseNodeClassificationTrainer,
+    BaseGraphClassificationTrainer,
+)
 
 
 def register_trainer(name):
@@ -19,9 +25,6 @@ def register_trainer(name):
     return register_trainer_cls
 
 
-EVALUATE_DICT = {}
-
-
 def register_evaluate(*name):
     def register_evaluate_cls(cls):
         for n in name:
@@ -36,6 +39,7 @@ def register_evaluate(*name):
 
     return register_evaluate_cls
 
+
 def get_feval(feval):
     if isinstance(feval, str):
         return EVALUATE_DICT[feval]
@@ -46,14 +50,17 @@ def get_feval(feval):
     raise ValueError("feval argument of type", type(feval), "is not supported!")
 
 
-from .graph_classification import GraphClassificationTrainer
-from .node_classification import NodeClassificationTrainer
+from .graph_classification_full import GraphClassificationFullTrainer
+from .node_classification_full import NodeClassificationFullTrainer
+from .node_classification_trainer import *
 from .evaluate import Acc, Auc, Logloss
 
 __all__ = [
     "BaseTrainer",
-    "GraphClassificationTrainer",
-    "NodeClassificationTrainer",
+    "BaseNodeClassificationTrainer",
+    "BaseGraphClassificationTrainer",
+    "GraphClassificationFullTrainer",
+    "NodeClassificationFullTrainer",
     "Evaluation",
     "Acc",
     "Auc",

autogl/module/train/base.py

@@ -1,12 +1,25 @@
 import numpy as np
 from typing import Union, Iterable
-from ..model import BaseModel
+
+import torch
+from ..model import BaseModel, MODEL_DICT
 import pickle
 from ...utils import get_logger
+from . import EVALUATE_DICT
 
 LOGGER_ES = get_logger("early-stopping")
 
 
+def get_feval(feval):
+    if isinstance(feval, str):
+        return EVALUATE_DICT[feval]
+    if isinstance(feval, type) and issubclass(feval, Evaluation):
+        return feval
+    if isinstance(feval, list):
+        return [get_feval(f) for f in feval]
+    raise ValueError("feval argument of type", type(feval), "is not supported!")
+
+
 class EarlyStopping:
     """Early stops the training if validation loss doesn't improve after a given patience."""
 
@@ -81,17 +94,11 @@ class EarlyStopping:
 class BaseTrainer:
     def __init__(
         self,
-        model: Union[BaseModel, str],
-        optimizer=None,
-        lr=None,
-        max_epoch=None,
-        early_stopping_round=None,
-        device=None,
+        model: BaseModel,
+        device: Union[torch.device, str],
         init=True,
         feval=["acc"],
         loss="nll_loss",
-        *args,
-        **kwargs,
     ):
         """
         The basic trainer.
@@ -103,29 +110,26 @@ class BaseTrainer:
         model: `BaseModel` or `str`
             The (name of) model used to train and predict.
 
-        optimizer: `Optimizer` of `str`
-            The (name of) optimizer used to train and predict.
-
-        lr: `float`
-            The learning rate.
-
-        max_epoch: `int`
-            The max number of epochs in training.
-
-        early_stopping_round: `int`
-            The round of early stop.
-
-        device: `torch.device` or `str`
-            The device where model will be running on.
-
         init: `bool`
             If True(False), the model will (not) be initialized.
+        """
+        super().__init__()
+        self.model = model
+        self.to(device)
+        self.init = init
+        self.feval = get_feval(feval)
+        self.loss = loss
 
-        args: Other parameters.
+    def to(self, device):
+        """
+        Migrate trainer to new device
 
-        kwargs: Other parameters.
+        Parameters
+        ----------
+        device: `str` or `torch.device`
+            The device this trainer will use
         """
-        super().__init__()
+        self.device = torch.device(device)
 
     def initialize(self):
         """Initialize the auto model in trainer."""
@@ -169,8 +173,8 @@ class BaseTrainer:
 
     @classmethod
     def load(cls, path):
-        with open(path, "rb") as input:
-            instance = pickle.load(input)
+        with open(path, "rb") as inputs:
+            instance = pickle.load(inputs)
             return instance
 
     @property
@@ -279,7 +283,21 @@ class BaseTrainer:
 
     def set_feval(self, feval):
         """Set the evaluation metrics."""
-        raise NotImplementedError()
+        self.feval = get_feval(feval)
+
+    def update_parameters(self, **kwargs):
+        """
+        Update parameters of this trainer
+        """
+        for k, v in kwargs.items():
+            if k == "feval":
+                self.set_feval(v)
+            elif k == "device":
+                self.to(v)
+            elif hasattr(self, k):
+                setattr(self, k, v)
+            else:
+                raise KeyError("Cannot set parameter", k, "for trainer", self.__class__)
 
 
 # a static class for evaluating results
@@ -296,7 +314,7 @@ class Evaluation:
         """
         Should return whether this evaluation method is higher better (bool)
         """
-        raise True
+        return True
 
     @staticmethod
     def evaluate(predict, label):
@@ -304,3 +322,84 @@ class Evaluation:
         Should return: the evaluation result (float)
         """
         raise NotImplementedError()
+
+
+class BaseNodeClassificationTrainer(BaseTrainer):
+    def __init__(
+        self,
+        model: Union[BaseModel, str],
+        num_features,
+        num_classes,
+        device="auto",
+        init=True,
+        feval=["acc"],
+        loss="nll_loss",
+    ):
+        self.num_features = num_features
+        self.num_classes = num_classes
+        device = (
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if device == "auto"
+            else torch.device(device)
+        )
+        if isinstance(model, str):
+            assert model in MODEL_DICT, "Cannot parse model name " + model
+            self.model = MODEL_DICT[model](num_features, num_classes, device, init=init)
+        elif isinstance(model, BaseModel):
+            self.model = model
+        else:
+            raise TypeError(
+                "Model argument only support str or BaseModel, get",
+                type(model),
+                "instead.",
+            )
+        super().__init__(model, device=device, init=init, feval=feval, loss=loss)
+
+    @classmethod
+    def get_task_name(cls):
+        return "GraphClassification"
+
+
+class BaseGraphClassificationTrainer(BaseTrainer):
+    def __init__(
+        self,
+        model: Union[BaseModel, str],
+        num_features,
+        num_classes,
+        num_graph_features=0,
+        device=None,
+        init=True,
+        feval=["acc"],
+        loss="nll_loss",
+    ):
+        self.num_features = num_features
+        self.num_classes = num_classes
+        self.num_graph_features = num_graph_features
+        device = (
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if device == "auto"
+            else torch.device(device)
+        )
+        if isinstance(model, str):
+            assert model in MODEL_DICT, "Cannot parse model name " + model
+            self.model = MODEL_DICT[model](
+                num_features,
+                num_classes,
+                device,
+                init=init,
+                num_graph_features=num_graph_features,
+            )
+        elif isinstance(model, BaseModel):
+            self.model = model
+        else:
+            raise TypeError(
+                "Model argument only support str or BaseModel, get",
+                type(model),
+                "instead.",
+            )
+
+        super().__init__(model, device=device, init=init, feval=feval, loss=loss)
+
+    @classmethod
+    def get_task_name(cls):
+        return "NodeClassification"

autogl/module/train/graph_classification.py → autogl/module/train/graph_classification_full.py

@@ -1,6 +1,12 @@
-from . import register_trainer, BaseTrainer, Evaluation, EVALUATE_DICT, EarlyStopping
+from . import register_trainer, EVALUATE_DICT
+from .base import BaseGraphClassificationTrainer, EarlyStopping, Evaluation
 import torch
-from torch.optim.lr_scheduler import StepLR
+from torch.optim.lr_scheduler import (
+    StepLR,
+    MultiStepLR,
+    ExponentialLR,
+    ReduceLROnPlateau,
+)
 import torch.nn.functional as F
 from ..model import MODEL_DICT, BaseModel
 from .evaluate import Logloss
@@ -11,7 +17,8 @@ import torch.multiprocessing as mp
 
 from ...utils import get_logger
 
-LOGGER = get_logger('graph classification solver')
+LOGGER = get_logger("graph classification solver")
+
 
 def get_feval(feval):
     if isinstance(feval, str):
@@ -23,8 +30,8 @@ def get_feval(feval):
     raise ValueError("feval argument of type", type(feval), "is not supported!")
 
 
-@register_trainer("GraphClassification")
-class GraphClassificationTrainer(BaseTrainer):
+@register_trainer("GraphClassificationFull")
+class GraphClassificationFullTrainer(BaseGraphClassificationTrainer):
     """
     The graph classification trainer.
 
@@ -69,30 +76,26 @@ class GraphClassificationTrainer(BaseTrainer):
         num_workers=None,
         early_stopping_round=7,
         weight_decay=1e-4,
-        device=None,
+        device="auto",
         init=True,
         feval=[Logloss],
         loss="nll_loss",
+        lr_scheduler_type=None,
         *args,
         **kwargs
     ):
-        super(GraphClassificationTrainer, self).__init__(model)
-
-        self.loss_type = loss
-
-        # init model
-        if isinstance(model, str):
-            assert model in MODEL_DICT, "Cannot parse model name " + model
-            self.model = MODEL_DICT[model](
-                num_features,
-                num_classes,
-                device,
-                init=init,
-                num_graph_features=num_graph_features,
-            )
-        elif isinstance(model, BaseModel):
-            self.model = model
+        super().__init__(
+            model,
+            num_features,
+            num_classes,
+            num_graph_features=num_graph_features,
+            device=device,
+            init=init,
+            feval=feval,
+            loss=loss,
+        )
 
+        self.opt_received = optimizer
         if type(optimizer) == str and optimizer.lower() == "adam":
             self.optimizer = torch.optim.Adam
         elif type(optimizer) == str and optimizer.lower() == "sgd":
@@ -100,9 +103,8 @@ class GraphClassificationTrainer(BaseTrainer):
         else:
             self.optimizer = torch.optim.Adam
 
-        self.num_features = num_features
-        self.num_classes = num_classes
-        self.num_graph_features = num_graph_features
+        self.lr_scheduler_type = lr_scheduler_type
+
         self.lr = lr if lr is not None else 1e-4
         self.max_epoch = max_epoch if max_epoch is not None else 100
         self.batch_size = batch_size if batch_size is not None else 64
@@ -130,8 +132,6 @@ class GraphClassificationTrainer(BaseTrainer):
         self.valid_score = None
 
         self.initialized = False
-        self.num_features = num_features
-        self.num_classes = num_classes
         self.device = device
 
         self.space = [
@@ -171,8 +171,6 @@ class GraphClassificationTrainer(BaseTrainer):
                 "scalingType": "LOG",
             },
         ]
-        self.space += self.model.space
-        GraphClassificationTrainer.space = self.space
 
         self.hyperparams = {
             "max_epoch": self.max_epoch,
@@ -181,7 +179,6 @@ class GraphClassificationTrainer(BaseTrainer):
             "lr": self.lr,
             "weight_decay": self.weight_decay,
         }
-        self.hyperparams = {**self.hyperparams, **self.model.get_hyper_parameter()}
 
         if init is True:
             self.initialize()
@@ -202,9 +199,9 @@ class GraphClassificationTrainer(BaseTrainer):
         # """Get task name, i.e., `GraphClassification`."""
         return "GraphClassification"
 
-    def to(self, new_device):
-        assert isinstance(new_device, torch.device)
-        self.device = new_device
+    def to(self, device):
+        assert isinstance(device, torch.device)
+        self.device = device
         if self.model is not None:
             self.model.to(self.device)
 
@@ -226,7 +223,22 @@ class GraphClassificationTrainer(BaseTrainer):
         optimizer = self.optimizer(
             self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay
         )
-        scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
+
+        # scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
+        lr_scheduler_type = self.lr_scheduler_type
+        if type(lr_scheduler_type) == str and lr_scheduler_type == "steplr":
+            scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
+        elif type(lr_scheduler_type) == str and lr_scheduler_type == "multisteplr":
+            scheduler = MultiStepLR(optimizer, milestones=[30, 80], gamma=0.1)
+        elif type(lr_scheduler_type) == str and lr_scheduler_type == "exponentiallr":
+            scheduler = ExponentialLR(optimizer, gamma=0.1)
+        elif (
+            type(lr_scheduler_type) == str and lr_scheduler_type == "reducelronplateau"
+        ):
+            scheduler = ReduceLROnPlateau(optimizer, "min")
+        else:
+            scheduler = None
+
         for epoch in range(1, self.max_epoch):
             self.model.model.train()
             loss_all = 0
@@ -235,29 +247,33 @@ class GraphClassificationTrainer(BaseTrainer):
                 optimizer.zero_grad()
                 output = self.model.model(data)
                 # loss = F.nll_loss(output, data.y)
-                if hasattr(F, self.loss_type):
-                    loss = getattr(F, self.loss_type)(output, data.y)
+                if hasattr(F, self.loss):
+                    loss = getattr(F, self.loss)(output, data.y)
                 else:
-                    raise TypeError("PyTorch does not support loss type {}".format(self.loss_type))
+                    raise TypeError(
+                        "PyTorch does not support loss type {}".format(self.loss)
+                    )
                 loss.backward()
                 loss_all += data.num_graphs * loss.item()
                 optimizer.step()
-                scheduler.step()
-
+                if self.lr_scheduler_type:
+                    scheduler.step()
             # loss = loss_all / len(train_loader.dataset)
             # train_loss = self.evaluate(train_loader)
-            eval_func = (
-                self.feval if not isinstance(self.feval, list) else self.feval[0]
-            )
-            val_loss = self._evaluate(valid_loader, eval_func) if valid_loader else 0.0
-
-            if eval_func.is_higher_better():
-                val_loss = -val_loss
-            self.early_stopping(val_loss, self.model.model)
-            if self.early_stopping.early_stop:
-                LOGGER.debug("Early stopping at", epoch)
-                self.early_stopping.load_checkpoint(self.model.model)
-                break
+            if valid_loader is not None:
+                eval_func = (
+                    self.feval if not isinstance(self.feval, list) else self.feval[0]
+                )
+                val_loss = self._evaluate(valid_loader, eval_func)
+
+                if eval_func.is_higher_better():
+                    val_loss = -val_loss
+                self.early_stopping(val_loss, self.model.model)
+                if self.early_stopping.early_stop:
+                    LOGGER.debug("Early stopping at", epoch)
+                    break
+        if valid_loader is not None:
+            self.early_stopping.load_checkpoint(self.model.model)
 
     def predict_only(self, loader):
         """
@@ -537,7 +553,7 @@ class GraphClassificationTrainer(BaseTrainer):
             num_features=self.num_features,
             num_classes=self.num_classes,
             num_graph_features=self.num_graph_features,
-            optimizer=self.optimizer,
+            optimizer=self.opt_received,
             lr=hp["lr"],
             max_epoch=hp["max_epoch"],
             batch_size=hp["batch_size"],
@@ -545,6 +561,8 @@ class GraphClassificationTrainer(BaseTrainer):
             weight_decay=hp["weight_decay"],
             device=self.device,
             feval=self.feval,
+            loss=self.loss,
+            lr_scheduler_type=self.lr_scheduler_type,
             init=True,
             *self.args,
             **self.kwargs
@@ -565,7 +583,6 @@ class GraphClassificationTrainer(BaseTrainer):
     def hyper_parameter_space(self, space):
         # """Set the space of hyperparameter."""
         self.space = space
-        GraphClassificationTrainer.space = space
 
     def get_hyper_parameter(self):
         # """Get the hyperparameter in this trainer."""

autogl/module/train/node_classification.py → autogl/module/train/node_classification_full.py

@@ -1,6 +1,17 @@
-from . import register_trainer, BaseTrainer, Evaluation, EVALUATE_DICT, EarlyStopping
+"""
+Node classification Full Trainer Implementation
+"""
+
+from . import register_trainer, EVALUATE_DICT
+
+from .base import BaseNodeClassificationTrainer, EarlyStopping, Evaluation
 import torch
-from torch.optim.lr_scheduler import StepLR
+from torch.optim.lr_scheduler import (
+    StepLR,
+    MultiStepLR,
+    ExponentialLR,
+    ReduceLROnPlateau,
+)
 import torch.nn.functional as F
 from ..model import MODEL_DICT, BaseModel
 from .evaluate import Logloss, Acc, Auc
@@ -11,6 +22,7 @@ from ...utils import get_logger
 
 LOGGER = get_logger("node classification trainer")
 
+
 def get_feval(feval):
     if isinstance(feval, str):
         return EVALUATE_DICT[feval]
@@ -21,8 +33,8 @@ def get_feval(feval):
     raise ValueError("feval argument of type", type(feval), "is not supported!")
 
 
-@register_trainer("NodeClassification")
-class NodeClassificationTrainer(BaseTrainer):
+@register_trainer("NodeClassificationFull")
+class NodeClassificationFullTrainer(BaseNodeClassificationTrainer):
     """
     The node classification trainer.
 
@@ -52,8 +64,6 @@ class NodeClassificationTrainer(BaseTrainer):
         If True(False), the model will (not) be initialized.
     """
 
-    space = None
-
     def __init__(
         self,
         model: Union[BaseModel, str],
@@ -64,19 +74,23 @@ class NodeClassificationTrainer(BaseTrainer):
         max_epoch=None,
         early_stopping_round=None,
         weight_decay=1e-4,
-        device=None,
+        device="auto",
         init=True,
         feval=[Logloss],
         loss="nll_loss",
+        lr_scheduler_type=None,
         *args,
         **kwargs
     ):
-        super(NodeClassificationTrainer, self).__init__(model)
-
-        self.loss_type = loss
-
-        if device is None:
-            device = "cpu"
+        super().__init__(
+            model,
+            num_features,
+            num_classes,
+            device=device,
+            init=init,
+            feval=feval,
+            loss=loss,
+        )
 
         # init model
         if isinstance(model, str):
@@ -85,6 +99,7 @@ class NodeClassificationTrainer(BaseTrainer):
         elif isinstance(model, BaseModel):
             self.model = model
 
+        self.opt_received = optimizer
         if type(optimizer) == str and optimizer.lower() == "adam":
             self.optimizer = torch.optim.Adam
         elif type(optimizer) == str and optimizer.lower() == "sgd":
@@ -92,14 +107,13 @@ class NodeClassificationTrainer(BaseTrainer):
         else:
             self.optimizer = torch.optim.Adam
 
-        self.num_features = num_features
-        self.num_classes = num_classes
+        self.lr_scheduler_type = lr_scheduler_type
+
         self.lr = lr if lr is not None else 1e-4
         self.max_epoch = max_epoch if max_epoch is not None else 100
         self.early_stopping_round = (
             early_stopping_round if early_stopping_round is not None else 100
         )
-        self.device = device
         self.args = args
         self.kwargs = kwargs
 
@@ -116,9 +130,6 @@ class NodeClassificationTrainer(BaseTrainer):
         self.valid_score = None
 
         self.initialized = False
-        self.num_features = num_features
-        self.num_classes = num_classes
-        self.device = device
 
         self.space = [
             {
@@ -150,8 +161,6 @@ class NodeClassificationTrainer(BaseTrainer):
                 "scalingType": "LOG",
             },
         ]
-        self.space += self.model.space
-        NodeClassificationTrainer.space = self.space
 
         self.hyperparams = {
             "max_epoch": self.max_epoch,
@@ -159,7 +168,6 @@ class NodeClassificationTrainer(BaseTrainer):
             "lr": self.lr,
             "weight_decay": self.weight_decay,
         }
-        self.hyperparams = {**self.hyperparams, **self.model.get_hyper_parameter()}
 
         if init is True:
             self.initialize()
@@ -200,32 +208,51 @@ class NodeClassificationTrainer(BaseTrainer):
         optimizer = self.optimizer(
             self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay
         )
-        scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
+        # scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
+        lr_scheduler_type = self.lr_scheduler_type
+        if type(lr_scheduler_type) == str and lr_scheduler_type == "steplr":
+            scheduler = StepLR(optimizer, step_size=100, gamma=0.1)
+        elif type(lr_scheduler_type) == str and lr_scheduler_type == "multisteplr":
+            scheduler = MultiStepLR(optimizer, milestones=[30, 80], gamma=0.1)
+        elif type(lr_scheduler_type) == str and lr_scheduler_type == "exponentiallr":
+            scheduler = ExponentialLR(optimizer, gamma=0.1)
+        elif (
+            type(lr_scheduler_type) == str and lr_scheduler_type == "reducelronplateau"
+        ):
+            scheduler = ReduceLROnPlateau(optimizer, "min")
+        else:
+            scheduler = None
+
         for epoch in range(1, self.max_epoch):
             self.model.model.train()
             optimizer.zero_grad()
             res = self.model.model.forward(data)
-            if hasattr(F, self.loss_type):
-                loss = getattr(F, self.loss_type)(res[mask], data.y[mask])
+            if hasattr(F, self.loss):
+                loss = getattr(F, self.loss)(res[mask], data.y[mask])
             else:
-                raise TypeError("PyTorch does not support loss type {}".format(self.loss_type))
+                raise TypeError(
+                    "PyTorch does not support loss type {}".format(self.loss)
+                )
 
             loss.backward()
             optimizer.step()
-            scheduler.step()
-
-            if type(self.feval) is list:
-                feval = self.feval[0]
-            else:
-                feval = self.feval
-            val_loss = self.evaluate([data], mask=data.val_mask, feval=feval)
-            if feval.is_higher_better() is True:
-                val_loss = -val_loss
-            self.early_stopping(val_loss, self.model.model)
-            if self.early_stopping.early_stop:
-                LOGGER.debug("Early stopping at %d", epoch)
-                self.early_stopping.load_checkpoint(self.model.model)
-                break
+            if self.lr_scheduler_type:
+                scheduler.step()
+
+            if hasattr(data, "val_mask") and data.val_mask is not None:
+                if type(self.feval) is list:
+                    feval = self.feval[0]
+                else:
+                    feval = self.feval
+                val_loss = self.evaluate([data], mask=data.val_mask, feval=feval)
+                if feval.is_higher_better() is True:
+                    val_loss = -val_loss
+                self.early_stopping(val_loss, self.model.model)
+                if self.early_stopping.early_stop:
+                    LOGGER.debug("Early stopping at %d", epoch)
+                    break
+        if hasattr(data, "val_mask") and data.val_mask is not None:
+            self.early_stopping.load_checkpoint(self.model.model)
 
     def predict_only(self, data, test_mask=None):
         """
@@ -480,13 +507,15 @@ class NodeClassificationTrainer(BaseTrainer):
             model=model,
             num_features=self.num_features,
             num_classes=self.num_classes,
-            optimizer=self.optimizer,
+            optimizer=self.opt_received,
             lr=hp["lr"],
             max_epoch=hp["max_epoch"],
             early_stopping_round=hp["early_stopping_round"],
             device=self.device,
             weight_decay=hp["weight_decay"],
             feval=self.feval,
+            loss=self.loss,
+            lr_scheduler_type=self.lr_scheduler_type,
             init=True,
             *self.args,
             **self.kwargs
@@ -507,7 +536,6 @@ class NodeClassificationTrainer(BaseTrainer):
     def hyper_parameter_space(self, space):
         # """Set the space of hyperparameter."""
         self.space = space
-        NodeClassificationTrainer.space = space
 
     def get_hyper_parameter(self):
         # """Get the hyperparameter in this trainer."""

autogl/module/train/node_classification_trainer/__init__.py

@@ -0,0 +1 @@
+from .node_classification_sampled_trainer import *

autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py

@@ -0,0 +1,424 @@
+import torch
+import logging
+import typing as _typing
+from torch.nn import functional as F
+
+from .. import EVALUATE_DICT, register_trainer
+from ..base import BaseNodeClassificationTrainer, EarlyStopping, Evaluation
+from ..evaluate import Logloss
+from ..sampling.sampler.neighbor_sampler import NeighborSampler
+from ...model import BaseModel, ModelUniversalRegistry
+
+LOGGER: logging.Logger = logging.getLogger("Node classification sampling trainer")
+
+
+def get_feval(feval):
+    if isinstance(feval, str):
+        return EVALUATE_DICT[feval]
+    if isinstance(feval, type) and issubclass(feval, Evaluation):
+        return feval
+    if isinstance(feval, list):
+        return [get_feval(f) for f in feval]
+    raise ValueError("feval argument of type", type(feval), "is not supported!")
+
+
+@register_trainer("NodeClassificationNeighborSampling")
+class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer):
+    """
+    The node classification trainer
+    for automatically training the node classification tasks
+    with neighbour sampling
+    """
+    
+    def __init__(
+            self,
+            model: _typing.Union[BaseModel, str],
+            num_features: int,
+            num_classes: int,
+            optimizer: _typing.Union[
+                _typing.Type[torch.optim.Optimizer], str, None
+            ] = None,
+            lr: float = 1e-4,
+            max_epoch: int = 100,
+            early_stopping_round: int = 100,
+            weight_decay: float = 1e-4,
+            device: _typing.Optional[torch.device] = None,
+            init: bool = True,
+            feval: _typing.Union[
+                _typing.Sequence[str],
+                _typing.Sequence[_typing.Type[Evaluation]]
+            ] = (Logloss,),
+            loss: str = "nll_loss",
+            lr_scheduler_type: _typing.Optional[str] = None,
+            **kwargs
+    ) -> None:
+        
+        self._functional_loss_name: str = loss
+        if device is None:
+            device: torch.device = torch.device("cpu")
+        
+        if type(model) == str:
+            self._model: BaseModel = ModelUniversalRegistry.get_model(model)(
+                num_features, num_classes, device, init=init
+            )
+        elif isinstance(model, BaseModel):
+            self._model: BaseModel = model
+        else:
+            raise TypeError
+        
+        if isinstance(optimizer, type) and issubclass(optimizer, torch.optim.Optimizer):
+            self._optimizer_class: _typing.Type[torch.optim.Optimizer] = optimizer
+        elif type(optimizer) == str:
+            if optimizer.lower() == "adam":
+                self._optimizer_class: _typing.Type[torch.optim.Optimizer] = torch.optim.Adam
+            elif optimizer.lower() == "adam" + "w":
+                self._optimizer_class: _typing.Type[torch.optim.Optimizer] = torch.optim.AdamW
+            elif optimizer.lower() == "sgd":
+                self._optimizer_class: _typing.Type[torch.optim.Optimizer] = torch.optim.SGD
+            else:
+                self._optimizer_class: _typing.Type[torch.optim.Optimizer] = torch.optim.Adam
+        else:
+            self._optimizer_class: _typing.Type[torch.optim.Optimizer] = torch.optim.Adam
+        
+        self._num_features: int = num_features
+        self._num_classes: int = num_classes
+        self._learning_rate: float = lr if lr > 0 else 1e-4
+        self._lr_scheduler_type: _typing.Optional[str] = lr_scheduler_type
+        self._max_epoch: int = max_epoch if max_epoch > 0 else 1e2
+        self._device: torch.device = device
+        
+        self.__sampling_sizes: _typing.Sequence[int] = kwargs.get("sampling_sizes")
+        
+        self._feval: _typing.Sequence[_typing.Type[Evaluation]] = get_feval(list(feval))
+        self._weight_decay: float = weight_decay if weight_decay > 0 else 1e-4
+        early_stopping_round: int = early_stopping_round if early_stopping_round > 0 else 1e2
+        self._early_stopping = EarlyStopping(patience=early_stopping_round, verbose=False)
+
+        super(NodeClassificationNeighborSamplingTrainer, self).__init__(
+            model, num_features, num_classes,
+            device=device if device is not None else "auto",
+            init=init, loss=loss
+        )
+        
+        self._valid_result: torch.Tensor = torch.zeros(0)
+        self._valid_result_prob: torch.Tensor = torch.zeros(0)
+        self._valid_score = None
+        
+        self._hyper_parameter_space: _typing.List[_typing.Dict[str, _typing.Any]] = [
+            {
+                "parameterName": "max_epoch",
+                "type": "INTEGER",
+                "maxValue": 500,
+                "minValue": 10,
+                "scalingType": "LINEAR",
+            },
+            {
+                "parameterName": "early_stopping_round",
+                "type": "INTEGER",
+                "maxValue": 30,
+                "minValue": 10,
+                "scalingType": "LINEAR",
+            },
+            {
+                "parameterName": "lr",
+                "type": "DOUBLE",
+                "maxValue": 1e-1,
+                "minValue": 1e-4,
+                "scalingType": "LOG",
+            },
+            {
+                "parameterName": "weight_decay",
+                "type": "DOUBLE",
+                "maxValue": 1e-2,
+                "minValue": 1e-4,
+                "scalingType": "LOG",
+            }
+        ]
+        
+        self._hyper_parameter: _typing.Dict[str, _typing.Any] = {
+            "max_epoch": self._max_epoch,
+            "early_stopping_round": self._early_stopping.patience,
+            "lr": self._learning_rate,
+            "weight_decay": self._weight_decay
+        }
+        
+        self.__initialized: bool = False
+        if init:
+            self.initialize()
+    
+    def initialize(self) -> "NodeClassificationNeighborSamplingTrainer":
+        if self.__initialized:
+            return self
+        self._model.initialize()
+        self.__initialized = True
+        return self
+    
+    def get_model(self) -> BaseModel:
+        return self._model
+    
+    def __train_only(
+            self, data
+    ) -> "NodeClassificationNeighborSamplingTrainer":
+        """
+        The function of training on the given dataset and mask.
+        :param data: data of a specific graph
+        :return: self
+        """
+        data = data.to(self._device)
+        optimizer: torch.optim.Optimizer = self._optimizer_class(
+            self._model.parameters(),
+            lr=self._learning_rate, weight_decay=self._weight_decay
+        )
+        if type(self._lr_scheduler_type) == str:
+            if self._lr_scheduler_type.lower() == "step" + "lr":
+                lr_scheduler: torch.optim.lr_scheduler.StepLR = \
+                    torch.optim.lr_scheduler.StepLR(
+                        optimizer, step_size=100, gamma=0.1
+                    )
+            elif self._lr_scheduler_type.lower() == "multi" + "step" + "lr":
+                lr_scheduler: torch.optim.lr_scheduler.MultiStepLR = \
+                    torch.optim.lr_scheduler.MultiStepLR(
+                        optimizer, milestones=[30, 80], gamma=0.1
+                    )
+            elif self._lr_scheduler_type.lower() == "exponential" + "lr":
+                lr_scheduler: torch.optim.lr_scheduler.ExponentialLR = \
+                    torch.optim.lr_scheduler.ExponentialLR(
+                        optimizer, gamma=0.1
+                    )
+            elif self._lr_scheduler_type.lower() == "ReduceLROnPlateau".lower():
+                lr_scheduler: torch.optim.lr_scheduler.ReduceLROnPlateau = \
+                    torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min")
+            else:
+                lr_scheduler: torch.optim.lr_scheduler.LambdaLR = \
+                    torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0)
+        else:
+            lr_scheduler: torch.optim.lr_scheduler.LambdaLR = \
+                torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0)
+        
+        train_sampler: NeighborSampler = NeighborSampler(
+            data, self.__sampling_sizes, batch_size=20
+        )
+        
+        for current_epoch in range(self._max_epoch):
+            self._model.model.train()
+            """ epoch start """
+            for target_node_indexes, edge_indexes in train_sampler:
+                optimizer.zero_grad()
+                data.edge_indexes = edge_indexes
+                prediction = self._model.model(data)
+                if not hasattr(F, self._functional_loss_name):
+                    raise TypeError(
+                        "PyTorch does not support loss type {}".format(self._functional_loss_name)
+                    )
+                loss_function = getattr(F, self._functional_loss_name)
+                loss: torch.Tensor = loss_function(
+                    prediction[target_node_indexes],
+                    data.y[target_node_indexes]
+                )
+                loss.backward()
+                optimizer.step()
+            
+            if lr_scheduler is not None:
+                lr_scheduler.step()
+            
+            """ Validate performance """
+            if hasattr(data, "val_mask") and getattr(data, "val_mask") is not None:
+                validation_results: _typing.Sequence[float] = \
+                    self.evaluate((data,), "val", [self._feval[0]])
+                
+                if self._feval[0].is_higher_better():
+                    validation_loss: float = -validation_results[0]
+                else:
+                    validation_loss: float = validation_results[0]
+                self._early_stopping(validation_loss, self._model.model)
+                if self._early_stopping.early_stop:
+                    LOGGER.debug("Early stopping at %d", current_epoch)
+                    break
+        if hasattr(data, "val_mask") and data.val_mask is not None:
+            self._early_stopping.load_checkpoint(self._model.model)
+        return self
+    
+    def __predict_only(self, data):
+        """
+        The function of predicting on the given data.
+        :param data: data of a specific graph
+        :return: the result of prediction on the given dataset
+        """
+        data = data.to(self._device)
+        self._model.model.eval()
+        with torch.no_grad():
+            prediction = self._model.model(data)
+        return prediction
+    
+    def train(self, dataset, keep_valid_result: bool = True):
+        """
+        The function of training on the given dataset and keeping valid result.
+        :param dataset:
+        :param keep_valid_result: Whether to save the validation result after training
+        """
+        data = dataset[0]
+        self.__train_only(data)
+        if keep_valid_result:
+            prediction: torch.Tensor = self.__predict_only(data)
+            self._valid_result: torch.Tensor = prediction[data.val_mask].max(1)[1]
+            self._valid_result_prob: torch.Tensor = prediction[data.val_mask]
+            self._valid_score = self.evaluate(dataset, "val")
+    
+    def predict_proba(
+            self, dataset, mask: _typing.Optional[str] = None,
+            in_log_format: bool = False
+    ) -> torch.Tensor:
+        """
+        The function of predicting the probability on the given dataset.
+        :param dataset: The node classification dataset used to be predicted.
+        :param mask:
+        :param in_log_format:
+        :return:
+        """
+        data = dataset[0].to(self._device)
+        if mask is not None and type(mask) == str:
+            if mask.lower() == "train":
+                _mask = data.train_mask
+            elif mask.lower() == "test":
+                _mask = data.test_mask
+            elif mask.lower() == "val":
+                _mask = data.val_mask
+            else:
+                _mask = data.test_mask
+        else:
+            _mask = data.test_mask
+        result = self.__predict_only(data)[_mask]
+        return result if in_log_format else torch.exp(result)
+    
+    def predict(self, dataset, mask: _typing.Optional[str] = None) -> torch.Tensor:
+        return self.predict_proba(
+            dataset, mask, in_log_format=True
+        ).max(1)[1]
+    
+    def get_valid_predict(self) -> torch.Tensor:
+        return self._valid_result
+    
+    def get_valid_predict_proba(self) -> torch.Tensor:
+        return self._valid_result_prob
+    
+    def get_valid_score(self, return_major: bool = True):
+        if return_major:
+            return (
+                self._valid_score[0],
+                self._feval[0].is_higher_better()
+            )
+        else:
+            return (
+                self._valid_score,
+                [f.is_higher_better() for f in self._feval]
+            )
+        
+    def get_name_with_hp(self) -> str:
+        # """Get the name of hyperparameter."""
+        name = "-".join(
+            [
+                str(self._optimizer_class),
+                str(self._learning_rate),
+                str(self._max_epoch),
+                str(self._early_stopping.patience),
+                str(self._model),
+                str(self._device),
+            ]
+        )
+        name = (
+            name
+            + "|"
+            + "-".join(
+                [
+                    str(x[0]) + "-" + str(x[1])
+                    for x in self.model.get_hyper_parameter().items()
+                ]
+            )
+        )
+        return name
+    
+    def evaluate(
+            self,
+            dataset,
+            mask: _typing.Optional[str] = None,
+            feval: _typing.Union[
+                None, _typing.Sequence[str],
+                _typing.Sequence[_typing.Type[Evaluation]]
+            ] = None
+    ) -> _typing.Sequence[float]:
+        data = dataset[0]
+        data = data.to(self._device)
+        if feval is None:
+            _feval: _typing.Sequence[_typing.Type[Evaluation]] = self._feval
+        else:
+            _feval: _typing.Sequence[_typing.Type[Evaluation]] = get_feval(list(feval))
+        if mask.lower() == "train":
+            _mask = data.train_mask
+        elif mask.lower() == "test":
+            _mask = data.test_mask
+        elif mask.lower() == "val":
+            _mask = data.val_mask
+        else:
+            _mask = data.test_mask
+        prediction_probability: torch.Tensor = self.predict_proba(dataset, mask)
+        y_ground_truth = data.y[_mask]
+        
+        results = []
+        for f in _feval:
+            try:
+                results.append(
+                    f.evaluate(prediction_probability, y_ground_truth)
+                )
+            except:
+                results.append(
+                    f.evaluate(prediction_probability.cpu().numpy(), y_ground_truth.cpu().numpy())
+                )
+        return results
+    
+    def to(self, device: torch.device):
+        self._device = device
+        if self._model is not None:
+            self._model.to(device)
+    
+    def duplicate_from_hyper_parameter(
+            self, hp: _typing.Dict[str, _typing.Any],
+            model: _typing.Union[BaseModel, str, None] = None
+    ) -> "NodeClassificationNeighborSamplingTrainer":
+        
+        if model is None or not isinstance(model, BaseModel):
+            model = self._model
+        model = model.from_hyper_parameter(
+            dict(
+                [
+                    x for x in hp.items()
+                    if x[0] in [y["parameterName"] for y in model.hyper_parameter_space]
+                ]
+            )
+        )
+        
+        return NodeClassificationNeighborSamplingTrainer(
+            model, self._num_features, self._num_classes,
+            self._optimizer_class,
+            device=self._device,
+            init=True,
+            feval=self._feval,
+            loss=self._functional_loss_name,
+            lr_scheduler_type=self._lr_scheduler_type,
+            **hp
+        )
+    
+    def set_feval(
+            self, feval: _typing.Union[
+                _typing.Sequence[str],
+                _typing.Sequence[_typing.Type[Evaluation]]
+            ]
+    ):
+        self._feval = get_feval(list(feval))
+    
+    @property
+    def hyper_parameter_space(self):
+        return self._hyper_parameter_space
+    
+    @hyper_parameter_space.setter
+    def hyper_parameter_space(self, hp_space):
+        self._hyper_parameter_space = hp_space

autogl/module/train/sampling/sampler/neighbor_sampler.py

@@ -0,0 +1,113 @@
+import collections
+import random
+import typing as _typing
+import numpy as np
+import torch.utils.data
+
+
+class NeighborSampler(torch.utils.data.DataLoader, collections.Iterable):
+    class _NodeIndexesDataset(torch.utils.data.Dataset):
+        def __init__(self, node_indexes):
+            self.__node_indexes: _typing.Sequence[int] = node_indexes
+        
+        def __getitem__(self, index) -> int:
+            if not 0 <= index < len(self.__node_indexes):
+                raise IndexError("Index out of range")
+            else:
+                return self.__node_indexes[index]
+        
+        def __len__(self) -> int:
+            return len(self.__node_indexes)
+    
+    def __init__(
+            self, data,
+            sampling_sizes: _typing.Sequence[int],
+            target_node_indexes: _typing.Optional[_typing.Sequence[int]] = None,
+            batch_size: _typing.Optional[int] = 1,
+            *args, **kwargs
+    ):
+        self._data = data
+        self.__sampling_sizes: _typing.Sequence[int] = sampling_sizes
+        
+        if not (
+                target_node_indexes is not None and
+                isinstance(target_node_indexes, _typing.Sequence)
+        ):
+            if hasattr(data, "train_mask"):
+                target_node_indexes: _typing.Sequence[int] = \
+                    torch.where(getattr(data, "train_mask"))[0]
+            else:
+                target_node_indexes: _typing.Sequence[int] = \
+                    list(np.arange(0, data.x.shape[0]))
+        
+        self.__edge_index_map: _typing.Dict[
+            int, _typing.Union[torch.Tensor, _typing.Sequence[int]]
+        ] = {}
+        self.__init_edge_index_map()
+        super(NeighborSampler, self).__init__(
+            self._NodeIndexesDataset(target_node_indexes),
+            batch_size=batch_size if batch_size > 0 else 1,
+            collate_fn=self.__sample, *args, **kwargs
+        )
+    
+    def __init_edge_index_map(self):
+        self.__edge_index_map.clear()
+        all_edge_index: torch.Tensor = getattr(self._data, "edge_index")
+        target_node_indexes: torch.Tensor = all_edge_index[1]
+        for target_node_index in target_node_indexes.unique().tolist():
+            self.__edge_index_map[target_node_index] = torch.where(
+                all_edge_index[1] == target_node_index
+            )[0]
+    
+    def __iter__(self):
+        return super(NeighborSampler, self).__iter__()
+    
+    def __sample(
+            self, target_nodes_indexes: _typing.List[int]
+    ) -> _typing.Tuple[torch.Tensor, _typing.List[torch.Tensor]]:
+        """
+        Sample a sub-graph with neighborhood sampling
+        :param target_nodes_indexes:
+        """
+        original_edge_index: torch.Tensor = self._data.edge_index
+        edges_indexes: _typing.List[torch.Tensor] = []
+        
+        current_target_nodes_indexes: _typing.List[int] = target_nodes_indexes
+        for current_sampling_size in self.__sampling_sizes:
+            current_edge_index: _typing.Optional[torch.Tensor] = None
+            for current_target_node_index in current_target_nodes_indexes:
+                if current_target_node_index in self.__edge_index_map:
+                    all_indexes: torch.Tensor = \
+                        self.__edge_index_map.get(current_target_node_index)
+                else:
+                    all_indexes: torch.Tensor = torch.where(
+                        original_edge_index[1] == current_target_node_index
+                    )[0]
+                if all_indexes.numel() < current_sampling_size:
+                    sampled_indexes: np.ndarray = np.random.choice(
+                        all_indexes.cpu().numpy(), current_sampling_size
+                    )
+                    if current_edge_index is not None:
+                        current_edge_index: torch.Tensor = torch.cat(
+                            [current_edge_index, original_edge_index[:, sampled_indexes]], dim=1
+                        )
+                    else:
+                        current_edge_index: torch.Tensor = original_edge_index[:, sampled_indexes]
+                else:
+                    all_indexes_list = all_indexes.tolist()
+                    random.shuffle(all_indexes_list)
+                    shuffled_indexes_list: _typing.List[int] = \
+                        all_indexes_list[0: current_sampling_size]
+                    if current_edge_index is not None:
+                        current_edge_index: torch.Tensor = torch.cat(
+                            [current_edge_index, original_edge_index[:, shuffled_indexes_list]], dim=1
+                        )
+                    else:
+                        current_edge_index: torch.Tensor = original_edge_index[:, shuffled_indexes_list]
+            edges_indexes.append(current_edge_index)
+            
+            if len(edges_indexes) < len(self.__sampling_sizes):
+                next_target_nodes_indexes: torch.Tensor = current_edge_index[0].unique()
+                current_target_nodes_indexes = next_target_nodes_indexes.tolist()
+        
+        return torch.tensor(target_nodes_indexes), edges_indexes[::-1]

autogl/solver/base.py

@@ -10,7 +10,7 @@ import torch
 
 from ..module.feature import FEATURE_DICT
 from ..module.hpo import HPO_DICT
-from ..module.train import NodeClassificationTrainer
+from ..module.model import MODEL_DICT
 from ..module import BaseFeatureAtom, BaseHPOptimizer, BaseTrainer
 from .utils import Leaderboard
 from ..utils import get_logger
@@ -18,8 +18,23 @@ from ..utils import get_logger
 LOGGER = get_logger("BaseSolver")
 
 
+def _initialize_single_model(model_name, parameters=None):
+    if parameters:
+        return MODEL_DICT[model_name](**parameters)
+    return MODEL_DICT[model_name]()
+
+
+def _parse_hp_space(spaces):
+    if spaces is None:
+        return None
+    for space in spaces:
+        if "cutFunc" in space and isinstance(space["cutFunc"], str):
+            space["cutFunc"] = eval(space["cutFunc"])
+    return spaces
+
+
 class BaseSolver:
-    """
+    r"""
     Base solver class, define some standard solver interfaces.
 
     Parameters
@@ -43,6 +58,12 @@ class BaseSolver:
         If given, will set the number eval times the hpo module will use.
         Only be effective when hpo_module is  of type ``str``. Default ``50``.
 
+    default_trainer: str or list of str (Optional)
+        Default trainer class to be used.
+        If a single trainer class is given, will set all trainer to default trainer.
+        If a list of trainer class is given, will set every model with corresponding trainer
+        cls. Default ``None``.
+
     trainer_hp_space: list of dict (Optional)
         trainer hp space or list of trainer hp spaces configuration.
         If a single trainer hp is given, will specify the hp space of trainer for every model.
@@ -71,6 +92,7 @@ class BaseSolver:
         hpo_module,
         ensemble_module,
         max_evals=50,
+        default_trainer=None,
         trainer_hp_space=None,
         model_hp_spaces=None,
         size=4,
@@ -87,12 +109,14 @@ class BaseSolver:
         elif isinstance(device, str) and (device == "cpu" or device.startswith("cuda")):
             self.runtime_device = torch.device(device)
         else:
-            LOGGER.error("Cannor parse device %s", str(device))
+            LOGGER.error("Cannot parse device %s", str(device))
             raise ValueError("Cannot parse device {}".format(device))
 
         # initialize modules
         self.graph_model_list = []
-        self.set_graph_models(graph_models, trainer_hp_space, model_hp_spaces)
+        self.set_graph_models(
+            graph_models, default_trainer, trainer_hp_space, model_hp_spaces
+        )
         self.set_feature_module(feature_module)
         self.set_hpo_module(hpo_module, max_evals=max_evals)
         self.set_ensemble_module(ensemble_module, size=size)
@@ -109,7 +133,7 @@ class BaseSolver:
         *args,
         **kwargs,
     ) -> "BaseSolver":
-        """
+        r"""
         Set the feature module of current solver.
 
         Parameters
@@ -159,10 +183,11 @@ class BaseSolver:
     def set_graph_models(
         self,
         graph_models,
+        default_trainer=None,
         trainer_hp_space=None,
         model_hp_spaces=None,
     ) -> "BaseSolver":
-        """
+        r"""
         Set the graph models used in current solver.
 
         Parameters
@@ -170,6 +195,12 @@ class BaseSolver:
         graph_models: list of autogl.module.model.BaseModel or list of str
             The (name of) models to be optimized as backbone.
 
+        default_trainer: str or list of str (Optional)
+            Default trainer class to be used.
+            If a single trainer class is given, will set all trainer to default trainer.
+            If a list of trainer class is given, will set every model with corresponding trainer
+            cls. Default ``None``.
+
         trainer_hp_space: list of dict (Optional)
             trainer hp space or list of trainer hp spaces configuration.
             If a single trainer hp is given, will specify the hp space of trainer for every model.
@@ -187,12 +218,13 @@ class BaseSolver:
             A reference of current solver.
         """
         self.gml = graph_models
+        self._default_trainer = default_trainer
         self._trainer_hp_space = trainer_hp_space
         self._model_hp_spaces = model_hp_spaces
         return self
 
     def set_hpo_module(self, hpo_module, *args, **kwargs) -> "BaseSolver":
-        """
+        r"""
         Set the hpo module used in current solver.
 
         Parameters
@@ -225,7 +257,7 @@ class BaseSolver:
             )
 
     def set_ensemble_module(self, ensemble_module, *args, **kwargs) -> "BaseSolver":
-        """
+        r"""
         Set the ensemble module used in current solver.
 
         Parameters
@@ -243,7 +275,7 @@ class BaseSolver:
         raise NotImplementedError()
 
     def fit(self, *args, **kwargs) -> "BaseSolver":
-        """
+        r"""
         Fit current solver on given dataset.
 
         Returns
@@ -254,7 +286,7 @@ class BaseSolver:
         raise NotImplementedError()
 
     def fit_predict(self, *args, **kwargs) -> Any:
-        """
+        r"""
         Fit current solver on given dataset and return the predicted value.
 
         Returns
@@ -265,7 +297,7 @@ class BaseSolver:
         raise NotImplementedError()
 
     def predict(self, *args, **kwargs) -> Any:
-        """
+        r"""
         Predict the node class number.
 
         Returns
@@ -276,7 +308,7 @@ class BaseSolver:
         raise NotImplementedError()
 
     def get_leaderboard(self) -> Leaderboard:
-        """
+        r"""
         Get the current leaderboard of this solver.
 
         Returns
@@ -287,7 +319,7 @@ class BaseSolver:
         return self.leaderboard
 
     def get_model_by_name(self, name) -> BaseTrainer:
-        """
+        r"""
         Find and get the model instance by name.
 
         Parameters
@@ -303,8 +335,8 @@ class BaseSolver:
         assert name in self.trained_models, "cannot find model by name" + name
         return self.trained_models[name]
 
-    def get_model_by_performance(self, index) -> Tuple[NodeClassificationTrainer, str]:
-        """
+    def get_model_by_performance(self, index) -> Tuple[BaseTrainer, str]:
+        r"""
         Find and get the model instance by performance.
 
         Parameters
@@ -314,7 +346,7 @@ class BaseSolver:
 
         Returns
         -------
-        trainer: autogl.module.train.NodeClassificationTrainer
+        trainer: autogl.module.train.BaseTrainer
             A trainer instance containing the trained models and training status.
         name: str
             The name of current trainer.
@@ -324,7 +356,7 @@ class BaseSolver:
 
     @classmethod
     def from_config(cls, path_or_dict, filetype="auto") -> "BaseSolver":
-        """
+        r"""
         Load solver from config file.
 
         You can use this function to directly load a solver from predefined config dict

autogl/solver/classifier/graph_classifier.py

@@ -12,9 +12,9 @@ import yaml
 
 from .base import BaseClassifier
 from ...module.feature import FEATURE_DICT
-from ...module.model import MODEL_DICT
-from ...module.train import TRAINER_DICT, get_feval
-from ...module import BaseModel
+from ...module.model import BaseModel, MODEL_DICT
+from ...module.train import TRAINER_DICT, get_feval, BaseGraphClassificationTrainer
+from ..base import _initialize_single_model, _parse_hp_space
 from ..utils import Leaderboard, set_seed
 from ...datasets import utils
 from ...utils import get_logger
@@ -77,6 +77,7 @@ class AutoGraphClassifier(BaseClassifier):
         hpo_module="anneal",
         ensemble_module="voting",
         max_evals=50,
+        default_trainer=None,
         trainer_hp_space=None,
         model_hp_spaces=None,
         size=4,
@@ -89,6 +90,7 @@ class AutoGraphClassifier(BaseClassifier):
             hpo_module=hpo_module,
             ensemble_module=ensemble_module,
             max_evals=max_evals,
+            default_trainer=default_trainer or "GraphClassificationFull",
             trainer_hp_space=trainer_hp_space,
             model_hp_spaces=model_hp_spaces,
             size=size,
@@ -100,10 +102,12 @@ class AutoGraphClassifier(BaseClassifier):
     def _init_graph_module(
         self,
         graph_models,
-        num_features,
         num_classes,
-        *args,
-        **kwargs,
+        num_features,
+        feval,
+        device,
+        loss,
+        num_graph_features,
     ) -> "AutoGraphClassifier":
         # load graph network module
         self.graph_model_list = []
@@ -113,10 +117,10 @@ class AutoGraphClassifier(BaseClassifier):
                     if model in MODEL_DICT:
                         self.graph_model_list.append(
                             MODEL_DICT[model](
-                                num_features=num_features,
                                 num_classes=num_classes,
-                                *args,
-                                **kwargs,
+                                num_features=num_features,
+                                num_graph_features=num_graph_features,
+                                device=device,
                                 init=False,
                             )
                         )
@@ -125,53 +129,79 @@ class AutoGraphClassifier(BaseClassifier):
                 elif isinstance(model, type) and issubclass(model, BaseModel):
                     self.graph_model_list.append(
                         model(
-                            num_features=num_features,
                             num_classes=num_classes,
-                            *args,
-                            **kwargs,
+                            num_features=num_features,
+                            num_graph_features=num_graph_features,
+                            device=device,
                             init=False,
                         )
                     )
                 elif isinstance(model, BaseModel):
-                    model.set_num_features(num_features)
+                    # setup the hp of num_classes and num_features
                     model.set_num_classes(num_classes)
-                    model.set_num_graph_features(
-                        0
-                        if "num_graph_features" not in kwargs
-                        else kwargs["num_graph_features"]
+                    model.set_num_features(num_features)
+                    model.set_num_graph_features(num_graph_features)
+                    self.graph_model_list.append(model.to(device))
+                elif isinstance(model, BaseGraphClassificationTrainer):
+                    # receive a trainer list, put trainer to list
+                    assert (
+                        model.get_model() is not None
+                    ), "Passed trainer should contain a model"
+                    model.model.set_num_classes(num_classes)
+                    model.model.set_num_features(num_features)
+                    model.model.set_num_graph_features(num_graph_features)
+                    model.update_parameters(
+                        num_classes=num_classes,
+                        num_features=num_features,
+                        num_graph_features=num_graph_features,
+                        loss=loss,
+                        feval=feval,
+                        device=device,
                     )
                     self.graph_model_list.append(model)
                 else:
                     raise KeyError("cannot find graph network %s." % (model))
         else:
             raise ValueError(
-                "need graph network to be str or a BaseModel class/instance, get",
+                "need graph network to be (list of) str or a BaseModel class/instance, get",
                 graph_models,
                 "instead.",
             )
 
         # wrap all model_cls with specified trainer
         for i, model in enumerate(self.graph_model_list):
+            # set model hp space
             if self._model_hp_spaces is not None:
                 if self._model_hp_spaces[i] is not None:
-                    model.hyper_parameter_space = self._model_hp_spaces[i]
-            trainer = TRAINER_DICT["GraphClassification"](
-                model=model,
-                num_features=num_features,
-                num_classes=num_classes,
-                *args,
-                **kwargs,
-                init=False,
-            )
+                    if isinstance(model, BaseGraphClassificationTrainer):
+                        model.model.hyper_parameter_space = self._model_hp_spaces[i]
+                    else:
+                        model.hyper_parameter_space = self._model_hp_spaces[i]
+            # initialize trainer if needed
+            if isinstance(model, BaseModel):
+                name = (
+                    self._default_trainer
+                    if isinstance(self._default_trainer, str)
+                    else self._default_trainer[i]
+                )
+                model = TRAINER_DICT[name](
+                    model=model,
+                    num_features=num_features,
+                    num_classes=num_classes,
+                    loss=loss,
+                    feval=feval,
+                    device=device,
+                    num_graph_features=num_graph_features,
+                    init=False,
+                )
+            # set trainer hp space
             if self._trainer_hp_space is not None:
                 if isinstance(self._trainer_hp_space[0], list):
                     current_hp_for_trainer = self._trainer_hp_space[i]
                 else:
                     current_hp_for_trainer = self._trainer_hp_space
-                trainer.hyper_parameter_space = (
-                    current_hp_for_trainer + model.hyper_parameter_space
-                )
-            self.graph_model_list[i] = trainer
+                model.hyper_parameter_space = current_hp_for_trainer
+            self.graph_model_list[i] = model
 
         return self
 
@@ -183,7 +213,7 @@ class AutoGraphClassifier(BaseClassifier):
         inplace=False,
         train_split=None,
         val_split=None,
-        cross_validation=True,
+        cross_validation=False,
         cv_split=10,
         evaluation_method="infer",
         seed=None,
@@ -215,7 +245,7 @@ class AutoGraphClassifier(BaseClassifier):
             Default ``None``.
 
         cross_validation: bool
-            Whether to use cross validation to fit on train dataset. Default ``True``.
+            Whether to use cross validation to fit on train dataset. Default ``False``.
 
         cv_split: int
             The cross validation split number. Only be effective when ``cross_validation=True``.
@@ -266,7 +296,7 @@ class AutoGraphClassifier(BaseClassifier):
                 "Please manually pass train and val ratio."
             )
             LOGGER.info("Use the default train/val/test ratio in given dataset")
-            #if hasattr(dataset.train_split, "n_splits"):
+            # if hasattr(dataset.train_split, "n_splits"):
             #    cross_validation = True
 
         elif train_split is not None and val_split is not None:
@@ -700,7 +730,7 @@ class AutoGraphClassifier(BaseClassifier):
         )
         if isinstance(path_or_dict, str):
             if filetype == "auto":
-                if path_or_dict.endswith(".yaml"):
+                if path_or_dict.endswith(".yaml") or path_or_dict.endswith(".yml"):
                     filetype = "yaml"
                 elif path_or_dict.endswith(".json"):
                     filetype = "json"
@@ -723,7 +753,7 @@ class AutoGraphClassifier(BaseClassifier):
         # load the dictionary
         path_or_dict = deepcopy(path_or_dict)
         solver = cls(None, [], None, None)
-        fe_list = path_or_dict.pop("feature", [{"name": "deepgl"}])
+        fe_list = path_or_dict.pop("feature", None)
         if fe_list is not None:
             fe_list_ele = []
             for feature_engineer in fe_list:
@@ -733,33 +763,51 @@ class AutoGraphClassifier(BaseClassifier):
             if fe_list_ele != []:
                 solver.set_feature_module(fe_list_ele)
 
-        models = path_or_dict.pop("models", {"gcn": None, "gat": None})
-        model_list = list(models.keys())
-        model_hp_space = [models[m] for m in model_list]
-        trainer_space = path_or_dict.pop("trainer", None)
-
-        # parse lambda function
-        if model_hp_space:
-            for space in model_hp_space:
-                if space is not None:
-                    for keys in space:
-                        if "cutFunc" in keys and isinstance(keys["cutFunc"], str):
-                            keys["cutFunc"] = eval(keys["cutFunc"])
-
-        if trainer_space:
-            for space in trainer_space:
-                if (
-                    isinstance(space, dict)
-                    and "cutFunc" in space
-                    and isinstance(space["cutFunc"], str)
-                ):
-                    space["cutFunc"] = eval(space["cutFunc"])
-                elif space is not None:
-                    for keys in space:
-                        if "cutFunc" in keys and isinstance(keys["cutFunc"], str):
-                            keys["cutFunc"] = eval(keys["cutFunc"])
-
-        solver.set_graph_models(model_list, trainer_space, model_hp_space)
+        models = path_or_dict.pop("models", [{"name": "gin"}, {"name": "topkpool"}])
+        model_hp_space = [
+            _parse_hp_space(model.pop("hp_space", None)) for model in models
+        ]
+        model_list = [
+            _initialize_single_model(model.pop("name"), model) for model in models
+        ]
+
+        trainer = path_or_dict.pop("trainer", None)
+        default_trainer = "GraphClassificationFull"
+        trainer_space = None
+        if isinstance(trainer, dict):
+            # global default
+            default_trainer = trainer.pop("name", "GraphClassificationFull")
+            trainer_space = _parse_hp_space(trainer.pop("hp_space", None))
+            default_kwargs = {"num_features": None, "num_classes": None}
+            default_kwargs.update(trainer)
+            default_kwargs["init"] = False
+            for i in range(len(model_list)):
+                model = model_list[i]
+                trainer_wrapper = TRAINER_DICT[default_trainer](
+                    model=model, **default_kwargs
+                )
+                model_list[i] = trainer_wrapper
+        elif isinstance(trainer, list):
+            # sequential trainer definition
+            assert len(trainer) == len(
+                model_list
+            ), "The number of trainer and model does not match"
+            trainer_space = []
+            for i in range(len(model_list)):
+                train, model = trainer[i], model_list[i]
+                default_trainer = train.pop("name", "GraphClassificationFull")
+                trainer_space.append(_parse_hp_space(train.pop("hp_space", None)))
+                default_kwargs = {"num_features": None, "num_classes": None}
+                default_kwargs.update(train)
+                default_kwargs["init"] = False
+                trainer_wrap = TRAINER_DICT[default_trainer](
+                    model=model, **default_kwargs
+                )
+                model_list[i] = trainer_wrap
+
+        solver.set_graph_models(
+            model_list, default_trainer, trainer_space, model_hp_space
+        )
 
         hpo_dict = path_or_dict.pop("hpo", {"name": "anneal"})
         if hpo_dict is not None:

autogl/solver/classifier/node_classifier.py

@@ -11,10 +11,11 @@ import numpy as np
 import yaml
 
 from .base import BaseClassifier
+from ..base import _parse_hp_space, _initialize_single_model
 from ...module.feature import FEATURE_DICT
-from ...module.model import MODEL_DICT
-from ...module.train import TRAINER_DICT, get_feval
-from ...module import BaseModel
+from ...module.model import MODEL_DICT, BaseModel
+from ...module.train import TRAINER_DICT, BaseNodeClassificationTrainer
+from ...module.train import get_feval
 from ..utils import Leaderboard, set_seed
 from ...datasets import utils
 from ...utils import get_logger
@@ -73,11 +74,12 @@ class AutoNodeClassifier(BaseClassifier):
 
     def __init__(
         self,
-        feature_module="deepgl",
+        feature_module=None,
         graph_models=["gat", "gcn"],
         hpo_module="anneal",
         ensemble_module="voting",
         max_evals=50,
+        default_trainer=None,
         trainer_hp_space=None,
         model_hp_spaces=None,
         size=4,
@@ -90,6 +92,7 @@ class AutoNodeClassifier(BaseClassifier):
             hpo_module=hpo_module,
             ensemble_module=ensemble_module,
             max_evals=max_evals,
+            default_trainer=default_trainer or "NodeClassificationFull",
             trainer_hp_space=trainer_hp_space,
             model_hp_spaces=model_hp_spaces,
             size=size,
@@ -100,12 +103,7 @@ class AutoNodeClassifier(BaseClassifier):
         self.data = None
 
     def _init_graph_module(
-        self,
-        graph_models,
-        num_classes,
-        num_features,
-        *args,
-        **kwargs,
+        self, graph_models, num_classes, num_features, feval, device, loss
     ) -> "AutoNodeClassifier":
         # load graph network module
         self.graph_model_list = []
@@ -117,8 +115,7 @@ class AutoNodeClassifier(BaseClassifier):
                             MODEL_DICT[model](
                                 num_classes=num_classes,
                                 num_features=num_features,
-                                *args,
-                                **kwargs,
+                                device=device,
                                 init=False,
                             )
                         )
@@ -129,8 +126,7 @@ class AutoNodeClassifier(BaseClassifier):
                         model(
                             num_classes=num_classes,
                             num_features=num_features,
-                            *args,
-                            **kwargs,
+                            device=device,
                             init=False,
                         )
                     )
@@ -138,6 +134,21 @@ class AutoNodeClassifier(BaseClassifier):
                     # setup the hp of num_classes and num_features
                     model.set_num_classes(num_classes)
                     model.set_num_features(num_features)
+                    self.graph_model_list.append(model.to(device))
+                elif isinstance(model, BaseNodeClassificationTrainer):
+                    # receive a trainer list, put trainer to list
+                    assert (
+                        model.get_model() is not None
+                    ), "Passed trainer should contain a model"
+                    model.model.set_num_classes(num_classes)
+                    model.model.set_num_features(num_features)
+                    model.update_parameters(
+                        num_classes=num_classes,
+                        num_features=num_features,
+                        loss=loss,
+                        feval=feval,
+                        device=device,
+                    )
                     self.graph_model_list.append(model)
                 else:
                     raise KeyError("cannot find graph network %s." % (model))
@@ -150,26 +161,37 @@ class AutoNodeClassifier(BaseClassifier):
 
         # wrap all model_cls with specified trainer
         for i, model in enumerate(self.graph_model_list):
+            # set model hp space
             if self._model_hp_spaces is not None:
                 if self._model_hp_spaces[i] is not None:
-                    model.hyper_parameter_space = self._model_hp_spaces[i]
-            trainer = TRAINER_DICT["NodeClassification"](
-                model=model,
-                num_features=num_features,
-                num_classes=num_classes,
-                *args,
-                **kwargs,
-                init=False,
-            )
+                    if isinstance(model, BaseNodeClassificationTrainer):
+                        model.model.hyper_parameter_space = self._model_hp_spaces[i]
+                    else:
+                        model.hyper_parameter_space = self._model_hp_spaces[i]
+            # initialize trainer if needed
+            if isinstance(model, BaseModel):
+                name = (
+                    self._default_trainer
+                    if isinstance(self._default_trainer, str)
+                    else self._default_trainer[i]
+                )
+                model = TRAINER_DICT[name](
+                    model=model,
+                    num_features=num_features,
+                    num_classes=num_classes,
+                    loss=loss,
+                    feval=feval,
+                    device=device,
+                    init=False,
+                )
+            # set trainer hp space
             if self._trainer_hp_space is not None:
                 if isinstance(self._trainer_hp_space[0], list):
                     current_hp_for_trainer = self._trainer_hp_space[i]
                 else:
                     current_hp_for_trainer = self._trainer_hp_space
-                trainer.hyper_parameter_space = (
-                    current_hp_for_trainer + model.hyper_parameter_space
-                )
-            self.graph_model_list[i] = trainer
+                model.hyper_parameter_space = current_hp_for_trainer
+            self.graph_model_list[i] = model
 
         return self
 
@@ -628,7 +650,7 @@ class AutoNodeClassifier(BaseClassifier):
         )
         if isinstance(path_or_dict, str):
             if filetype == "auto":
-                if path_or_dict.endswith(".yaml"):
+                if path_or_dict.endswith(".yaml") or path_or_dict.endswith(".yml"):
                     filetype = "yaml"
                 elif path_or_dict.endswith(".json"):
                     filetype = "json"
@@ -650,7 +672,7 @@ class AutoNodeClassifier(BaseClassifier):
 
         path_or_dict = deepcopy(path_or_dict)
         solver = cls(None, [], None, None)
-        fe_list = path_or_dict.pop("feature", [{"name": "deepgl"}])
+        fe_list = path_or_dict.pop("feature", None)
         if fe_list is not None:
             fe_list_ele = []
             for feature_engineer in fe_list:
@@ -660,33 +682,51 @@ class AutoNodeClassifier(BaseClassifier):
             if fe_list_ele != []:
                 solver.set_feature_module(fe_list_ele)
 
-        models = path_or_dict.pop("models", {"gcn": None, "gat": None})
-        model_list = list(models.keys())
-        model_hp_space = [models[m] for m in model_list]
-        trainer_space = path_or_dict.pop("trainer", None)
-
-        if model_hp_space:
-            # parse lambda function
-            for space in model_hp_space:
-                if space is not None:
-                    for keys in space:
-                        if "cutFunc" in keys and isinstance(keys["cutFunc"], str):
-                            keys["cutFunc"] = eval(keys["cutFunc"])
-
-        if trainer_space:
-            for space in trainer_space:
-                if (
-                    isinstance(space, dict)
-                    and "cutFunc" in space
-                    and isinstance(space["cutFunc"], str)
-                ):
-                    space["cutFunc"] = eval(space["cutFunc"])
-                elif space is not None:
-                    for keys in space:
-                        if "cutFunc" in keys and isinstance(keys["cutFunc"], str):
-                            keys["cutFunc"] = eval(keys["cutFunc"])
-
-        solver.set_graph_models(model_list, trainer_space, model_hp_space)
+        models = path_or_dict.pop("models", [{"name": "gcn"}, {"name": "gat"}])
+        model_hp_space = [
+            _parse_hp_space(model.pop("hp_space", None)) for model in models
+        ]
+        model_list = [
+            _initialize_single_model(model.pop("name"), model) for model in models
+        ]
+
+        trainer = path_or_dict.pop("trainer", None)
+        default_trainer = "NodeClassificationFull"
+        trainer_space = None
+        if isinstance(trainer, dict):
+            # global default
+            default_trainer = trainer.pop("name", "NodeClassificationFull")
+            trainer_space = _parse_hp_space(trainer.pop("hp_space", None))
+            default_kwargs = {"num_features": None, "num_classes": None}
+            default_kwargs.update(trainer)
+            default_kwargs["init"] = False
+            for i in range(len(model_list)):
+                model = model_list[i]
+                trainer_wrap = TRAINER_DICT[default_trainer](
+                    model=model, **default_kwargs
+                )
+                model_list[i] = trainer_wrap
+        elif isinstance(trainer, list):
+            # sequential trainer definition
+            assert len(trainer) == len(
+                model_list
+            ), "The number of trainer and model does not match"
+            trainer_space = []
+            for i in range(len(model_list)):
+                train, model = trainer[i], model_list[i]
+                default_trainer = train.pop("name", "NodeClassificationFull")
+                trainer_space.append(_parse_hp_space(train.pop("hp_space", None)))
+                default_kwargs = {"num_features": None, "num_classes": None}
+                default_kwargs.update(train)
+                default_kwargs["init"] = False
+                trainer_wrap = TRAINER_DICT[default_trainer](
+                    model=model, **default_kwargs
+                )
+                model_list[i] = trainer_wrap
+
+        solver.set_graph_models(
+            model_list, default_trainer, trainer_space, model_hp_space
+        )
 
         hpo_dict = path_or_dict.pop("hpo", {"name": "anneal"})
         if hpo_dict is not None:

configs/gcl_gin.yaml

@@ -1,65 +0,0 @@
-feature:
-  - name : ~
-
-models:
-  gin:
-    - parameterName: num_layers
-      type: FIXED
-      value: 6
-    
-    - parameterName: hidden
-      type: FIXED
-      value: [32,32,32,32,32]
-    
-    - parameterName: dropout
-      type: FIXED
-      value: 0.5
-    
-    - parameterName: act
-      type: FIXED
-      value: relu
-
-    - parameterName: eps
-      type: FIXED
-      value: True
-
-    - parameterName: mlp_layers
-      type: FIXED
-      value: 2
-
-trainer:
-  - parameterName: max_epoch
-    type: FIXED
-    value: 350
-  
-  - parameterName: early_stopping_round
-    type: FIXED
-    value: 10
-
-  - parameterName: lr
-    type: FIXED
-    value: 0.01
-  
-  - parameterName: weight_decay
-    type: FIXED
-    value: 0
-
-  - parameterName: batch_size
-    type: FIXED
-    value: 32
-
-# hidden tuned in {16,32} for bioinformatics,64 for social
-# batch tuned in {32,128}
-# dropout tuned in {0,0.5}
-
-# weight decay （0.5 every 50 epochs)
-
-# max epoch 350 
-# early stop epochs (run to end?), best for 10 folds
-
-hpo:
-  name: random
-  max_evals: 1
-
-ensemble:
-  name: ~

configs/graph_classification.yaml

@@ -1,66 +0,0 @@
-feature:
-  - name: NxLargeCliqueSize
-  - name: NxLargeCliqueSize
-
-models:
-  topkpool:
-
-    - parameterName: ratio
-      type: DOUBLE
-      maxValue: 0.9
-      minValue: 0.1
-      scalingType: LINEAR
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.9
-      minValue: 0.1
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
-trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
-    minValue: 10
-    scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
-    type: INTEGER
-    maxValue: 30
-    minValue: 10
-    scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.1
-    minValue: 0.0001
-    scalingType: LOG
-  
-  - parameterName: weight_decay
-    type: DOUBLE
-    maxValue: 0.005
-    minValue: 0.00005
-    scalingType: LOG
-
-  - parameterName: batch_size
-    type: INTEGER
-    maxValue: 128
-    minValue: 48
-    scalingType: LINEAR
-
-
-hpo:
-  name: anneal
-  max_evals: 10
-
-ensemble:
-  name: voting
-  size: 2

configs/graphclf_full.yml

@@ -0,0 +1,53 @@
+ensemble:
+  name: voting
+  size: 2
+hpo:
+  max_evals: 10
+  name: anneal
+models:
+- hp_space:
+  - maxValue: 0.9
+    minValue: 0.1
+    parameterName: ratio
+    scalingType: LINEAR
+    type: DOUBLE
+  - maxValue: 0.9
+    minValue: 0.1
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: topkpool
+trainer:
+  hp_space:
+  - maxValue: 300
+    minValue: 10
+    parameterName: max_epoch
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 30
+    minValue: 10
+    parameterName: early_stopping_round
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 0.1
+    minValue: 0.0001
+    parameterName: lr
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 0.005
+    minValue: 5.0e-05
+    parameterName: weight_decay
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 128
+    minValue: 48
+    parameterName: batch_size
+    scalingType: LINEAR
+    type: INTEGER

configs/graphclf_gin_benchmark.yml

@@ -0,0 +1,70 @@
+hpo:
+  max_evals: 10
+  name: tpe
+models:
+- hp_space:
+  - parameterName: num_layers
+    type: DISCRETE
+    feasiblePoints: '3,4,5'
+
+  - parameterName: hidden
+    type: NUMERICAL_LIST
+    numericalType: INTEGER
+    length: 5
+    minValue: [8, 8, 8, 8, 8]
+    maxValue: [64, 64, 64, 64, 64]
+    scalingType: LOG
+    cutPara: ["num_layers"]
+    cutFunc: "lambda x: x[0] - 1"
+
+  - parameterName: dropout
+    type: DOUBLE
+    maxValue: 0.9
+    minValue: 0.1
+    scalingType: LINEAR
+
+  - parameterName: act
+    type: CATEGORICAL
+    feasiblePoints:
+      - leaky_relu
+      - relu
+      - elu
+      - tanh
+
+  - parameterName: eps
+    type: CATEGORICAL
+    feasiblePoints:
+      - True
+      - False
+
+  - parameterName: mlp_layers
+    type: DISCRETE
+    feasiblePoints: '2,3,4'
+  name: gin
+trainer:
+  hp_space:
+  - maxValue: 300
+    minValue: 10
+    parameterName: max_epoch
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 30
+    minValue: 10
+    parameterName: early_stopping_round
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 0.1
+    minValue: 0.0001
+    parameterName: lr
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 0.005
+    minValue: 5.0e-05
+    parameterName: weight_decay
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 128
+    minValue: 48
+    parameterName: batch_size
+    scalingType: LINEAR
+    type: INTEGER

configs/graphclf_topk_benchmark.yml

@@ -0,0 +1,50 @@
+hpo:
+  max_evals: 10
+  name: tpe
+models:
+- hp_space:
+  - maxValue: 0.9
+    minValue: 0.1
+    parameterName: ratio
+    scalingType: LINEAR
+    type: DOUBLE
+  - maxValue: 0.9
+    minValue: 0.1
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: topkpool
+trainer:
+  hp_space:
+  - maxValue: 300
+    minValue: 10
+    parameterName: max_epoch
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 30
+    minValue: 10
+    parameterName: early_stopping_round
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 0.1
+    minValue: 0.0001
+    parameterName: lr
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 0.005
+    minValue: 5.0e-05
+    parameterName: weight_decay
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 128
+    minValue: 48
+    parameterName: batch_size
+    scalingType: LINEAR
+    type: INTEGER

configs/ncl_gat.yaml

@@ -1,52 +0,0 @@
-feature:
-  - name: ~
-
-models:
-  gat:
-    - parameterName: num_layers
-      type: FIXED
-      value: 2
-
-    - parameterName: heads
-      type: FIXED
-      value: 8
-    
-    - parameterName: hidden
-      type: FIXED
-      value: [64]
-    
-    - parameterName: dropout
-      type: FIXED
-      value: 0.6
-    
-    - parameterName: act
-      type: FIXED
-      value: elu
-
-trainer:
-  - parameterName: max_epoch
-    type: FIXED
-    value: 200
-  
-  - parameterName: early_stopping_round
-    type: FIXED
-    value: 10
-
-  - parameterName: lr
-    type: FIXED
-    value: 0.005
-  
-  - parameterName: weight_decay
-    type: FIXED
-    value: 0.0005
-
-# Glorot initialization
-# for pumbed dataset , heads = 8 for last layer  and weight decay =0.001 ,lr=0.01
-# early stopping 100, max epoch 100000
-hpo:
-  name: random
-  max_evals: 1
-
-ensemble:
-  name: ~
-

configs/ncl_gcn.yaml

@@ -1,45 +0,0 @@
-feature:
-  - name: ~ # ~ means None
-
-models:
-  gcn:
-    - parameterName: num_layers
-      type: FIXED
-      value: 3
-    
-    - parameterName: hidden
-      type: FIXED
-      value: [16, 16]
-    
-    - parameterName: dropout
-      type: FIXED
-      value: 0.5
-    
-    - parameterName: act
-      type: FIXED
-      value: relu
-
-trainer:
-  - parameterName: max_epoch
-    type: FIXED
-    value: 200
-  
-  - parameterName: early_stopping_round
-    type: FIXED
-    value: 10
-
-  - parameterName: lr
-    type: FIXED
-    value: 0.01
-  
-  - parameterName: weight_decay
-    type: FIXED
-    value: 0.0005
-# Glorot initialization 
-# weight decay only for the first layer
-hpo:
-  name: random
-  max_evals: 1
-
-ensemble:
-  name: ~

configs/node_classification.yaml

@@ -1,70 +0,0 @@
-feature:
-  - name: PYGNormalizeFeatures
-  - name: pagerank
-
-models:
-  gat:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2,3,4'
-
-    - parameterName: heads
-      type: DISCRETE
-      feasiblePoints: '4,8,16'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 3
-      minValue: [8, 8, 8]
-      maxValue: [64, 64, 64]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.9
-      minValue: 0.1
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
-trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
-    minValue: 10
-    scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
-    type: INTEGER
-    maxValue: 30
-    minValue: 10
-    scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.1
-    minValue: 0.0001
-    scalingType: LOG
-  
-  - parameterName: weight_decay
-    type: DOUBLE
-    maxValue: 0.005
-    minValue: 0.00005
-    scalingType: LOG
-
-hpo:
-  name: anneal
-  max_evals: 10
-
-ensemble:
-  name: voting
-  size: 2

configs/nodeclf_full.yml

@@ -0,0 +1,93 @@
+ensemble:
+  name: voting
+  size: 2
+feature:
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 50
+  name: tpe
+models:
+- hp_space:
+  - feasiblePoints: '2'
+    parameterName: num_layers
+    type: DISCRETE
+  - feasiblePoints: 6,8,10,12
+    parameterName: heads
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 1
+    maxValue:
+    - 16
+    minValue:
+    - 4
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gat
+- hp_space:
+  - feasiblePoints: '2'
+    parameterName: num_layers
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 1
+    maxValue:
+    - 64
+    minValue:
+    - 16
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gcn
+trainer:
+  hp_space:
+  - maxValue: 300
+    minValue: 100
+    parameterName: max_epoch
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 30
+    minValue: 10
+    parameterName: early_stopping_round
+    scalingType: LINEAR
+    type: INTEGER
+  - maxValue: 0.05
+    minValue: 0.01
+    parameterName: lr
+    scalingType: LOG
+    type: DOUBLE
+  - maxValue: 0.001
+    minValue: 0.0001
+    parameterName: weight_decay
+    scalingType: LOG
+    type: DOUBLE
+

configs/nodeclf_gat_benchmark_large.yml

@@ -1,69 +1,64 @@
-# search space for gat on amazon_computers amazon_photo coauthor_cs coauthor_physics
+ensemble:
+  name: null
 feature:
-  - name: PYGNormalizeFeatures
-
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
 models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2,3'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 2
-      minValue: [8, 8]
-      maxValue: [32, 32]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.5
-      minValue: 0.2
-      scalingType: LINEAR
-
-    - parameterName: heads
-      type: DISCRETE
-      feasiblePoints: '8,10,12'
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
+- hp_space:
+  - feasiblePoints: 2,3
+    parameterName: num_layers
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 2
+    maxValue:
+    - 32
+    - 32
+    minValue:
+    - 8
+    - 8
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.5
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints: 8,10,12
+    parameterName: heads
+    type: DISCRETE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gcn
 trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 400
+  hp_space:
+  - maxValue: 400
     minValue: 250
+    parameterName: max_epoch
     scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
     type: INTEGER
-    maxValue: 40
+  - maxValue: 40
     minValue: 25
+    parameterName: early_stopping_round
     scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.05
+    type: INTEGER
+  - maxValue: 0.05
     minValue: 0.01
+    parameterName: lr
     scalingType: LOG
-  
-  - parameterName: weight_decay
     type: DOUBLE
-    maxValue: 0.0005
+  - maxValue: 0.0005
     minValue: 0.0001
+    parameterName: weight_decay
     scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~
+    type: DOUBLE

configs/nodeclf_gat_benchmark_small.yml

@@ -1,70 +1,62 @@
-# search space for gat on cora, citeseer, pubmed
+ensemble:
+  name: null
 feature:
-  - name: PYGNormalizeFeatures
-
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
 models:
-  gat:
-
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2'
-  
-    - parameterName: heads
-      type: DISCRETE
-      feasiblePoints: '6,8,10,12'
-
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 1
-      minValue: [4]
-      maxValue: [16]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
+- hp_space:
+  - feasiblePoints: '2'
+    parameterName: num_layers
+    type: DISCRETE
+  - feasiblePoints: 6,8,10,12
+    parameterName: heads
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 1
+    maxValue:
+    - 16
+    minValue:
+    - 4
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gat
 trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
+  hp_space:
+  - maxValue: 300
     minValue: 100
+    parameterName: max_epoch
     scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
     type: INTEGER
-    maxValue: 30
+  - maxValue: 30
     minValue: 10
+    parameterName: early_stopping_round
     scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.05
+    type: INTEGER
+  - maxValue: 0.05
     minValue: 0.01
+    parameterName: lr
     scalingType: LOG
-  
-  - parameterName: weight_decay
     type: DOUBLE
-    maxValue: 0.001
+  - maxValue: 0.001
     minValue: 0.0001
+    parameterName: weight_decay
     scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~
+    type: DOUBLE

configs/nodeclf_gcn.yaml

@@ -1,64 +0,0 @@
-feature:
-  - name: ~ # ~ means None
-
-models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 2
-      minValue: [16, 16]
-      maxValue: [64, 64]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
-trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
-    minValue: 100
-    scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
-    type: INTEGER
-    maxValue: 30
-    minValue: 10
-    scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.01
-    minValue: 0.0025
-    scalingType: LOG
-  
-  - parameterName: weight_decay
-    type: DOUBLE
-    maxValue: 0.025
-    minValue: 0.0025
-    scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~

configs/nodeclf_gcn_benchmark_large.yml

@@ -1,65 +1,61 @@
-# search space for gcn on amazon_computers amazon_photo coauthor_cs coauthor_physics
+ensemble:
+  name: null
 feature:
-  - name: PYGNormalizeFeatures
-
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
 models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2,3'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 2
-      minValue: [32, 32]
-      maxValue: [128, 128]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
+- hp_space:
+  - feasiblePoints: 2,3
+    parameterName: num_layers
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 2
+    maxValue:
+    - 128
+    - 128
+    minValue:
+    - 32
+    - 32
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gcn
 trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
+  hp_space:
+  - maxValue: 300
     minValue: 100
+    parameterName: max_epoch
     scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
     type: INTEGER
-    maxValue: 30
+  - maxValue: 30
     minValue: 10
+    parameterName: early_stopping_round
     scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.05
+    type: INTEGER
+  - maxValue: 0.05
     minValue: 0.01
+    parameterName: lr
     scalingType: LOG
-  
-  - parameterName: weight_decay
     type: DOUBLE
-    maxValue: 0.0005
-    minValue: 0.00005
+  - maxValue: 0.0005
+    minValue: 5.0e-05
+    parameterName: weight_decay
     scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~
+    type: DOUBLE

configs/nodeclf_gcn_benchmark_small.yml

@@ -1,65 +1,59 @@
-# search space for gcn on cora, citeseer, pubmed
+ensemble:
+  name: null
 feature:
-  - name: PYGNormalizeFeatures
-
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
 models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 1
-      minValue: [16]
-      maxValue: [64]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
+- hp_space:
+  - feasiblePoints: '2'
+    parameterName: num_layers
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 1
+    maxValue:
+    - 64
+    minValue:
+    - 16
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gcn
 trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
+  hp_space:
+  - maxValue: 300
     minValue: 100
+    parameterName: max_epoch
     scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
     type: INTEGER
-    maxValue: 30
+  - maxValue: 30
     minValue: 10
+    parameterName: early_stopping_round
     scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.05
+    type: INTEGER
+  - maxValue: 0.05
     minValue: 0.005
+    parameterName: lr
     scalingType: LOG
-  
-  - parameterName: weight_decay
     type: DOUBLE
-    maxValue: 0.001
+  - maxValue: 0.001
     minValue: 0.0001
+    parameterName: weight_decay
     scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~
+    type: DOUBLE

configs/nodeclf_gcn_large.yaml

@@ -1,64 +0,0 @@
-feature:
-  - name: ~ # ~ means None
-
-models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2,3'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 2
-      minValue: [32, 32]
-      maxValue: [128, 128]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
-trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
-    minValue: 100
-    scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
-    type: INTEGER
-    maxValue: 30
-    minValue: 10
-    scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.01
-    minValue: 0.001
-    scalingType: LOG
-  
-  - parameterName: weight_decay
-    type: DOUBLE
-    maxValue: 0.01
-    minValue: 0.001
-    scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~

configs/nodeclf_sage_benchmark_large.yml

@@ -1,70 +1,76 @@
-# search space for graphsage on amazon_computers amazon_photo coauthor_cs coauthor_physics
+ensemble:
+  name: null
 feature:
-  - name: PYGNormalizeFeatures
-
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
 models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2,3'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 2
-      minValue: [32,128]
-      maxValue: [32,128]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-
-    - parameterName: agg, 
-      type: CATEGORICAL,
-      feasiblePoints": 
-        - mean
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
+- hp_space:
+  - feasiblePoints: 2,3
+    parameterName: num_layers
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 2
+    maxValue:
+    - 32
+    - 128
+    minValue:
+    - 32
+    - 128
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints":
+    - mean
+    parameterName: aggr,
+    type: CATEGORICAL,
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: sage
 trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
+  name: NodeClassificationNeighborSampling
+  hp_space:
+  - parameterName: sampling_sizes
+    type: NUMERICAL_LIST
+    numericalType: INTEGER
+    length: 3
+    cutFunc: lambda x:x[0]
+    cutPara:
+    - num_layers
+    minValue: 3
+    maxValue: 8
+    scalingType: LOG
+  - maxValue: 300
     minValue: 100
+    parameterName: max_epoch
     scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
     type: INTEGER
-    maxValue: 30
+  - maxValue: 30
     minValue: 10
+    parameterName: early_stopping_round
     scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.05
+    type: INTEGER
+  - maxValue: 0.05
     minValue: 0.01
+    parameterName: lr
     scalingType: LOG
-  
-  - parameterName: weight_decay
     type: DOUBLE
-    maxValue: 0.0005
+  - maxValue: 0.0005
     minValue: 0.0001
+    parameterName: weight_decay
     scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~
+    type: DOUBLE

configs/nodeclf_sage_benchmark_small.yml

@@ -1,72 +1,67 @@
-# search space for graphsage on cora, citeseer, pubmed
+ensemble:
+  name: null
 feature:
-  - name: PYGNormalizeFeatures
-
+- name: PYGNormalizeFeatures
+hpo:
+  max_evals: 10
+  name: random
 models:
-  gcn:
-    - parameterName: num_layers
-      type: DISCRETE
-      feasiblePoints: '2,3'
-    
-    - parameterName: hidden
-      type: NUMERICAL_LIST
-      numericalType: INTEGER
-      length: 2
-      minValue: [16,64]
-      maxValue: [16,64]
-      cutPara: ["num_layers"]
-      cutFunc: "lambda x:x[0] - 1"
-      scalingType: LOG
-    
-    - parameterName: dropout
-      type: DOUBLE
-      maxValue: 0.8
-      minValue: 0.2
-      scalingType: LINEAR
-
-    - parameterName: agg, 
-      type: CATEGORICAL,
-      feasiblePoints": 
-        - mean
-        - add
-        - max
-    
-    - parameterName: act
-      type: CATEGORICAL
-      feasiblePoints:
-        - leaky_relu
-        - relu
-        - elu
-        - tanh
-
+- hp_space:
+  - feasiblePoints: 2,3
+    parameterName: num_layers
+    type: DISCRETE
+  - cutFunc: lambda x:x[0] - 1
+    cutPara:
+    - num_layers
+    length: 2
+    maxValue:
+    - 16
+    - 64
+    minValue:
+    - 16
+    - 64
+    numericalType: INTEGER
+    parameterName: hidden
+    scalingType: LOG
+    type: NUMERICAL_LIST
+  - maxValue: 0.8
+    minValue: 0.2
+    parameterName: dropout
+    scalingType: LINEAR
+    type: DOUBLE
+  - feasiblePoints":
+    - mean
+    - add
+    - max
+    parameterName: agg,
+    type: CATEGORICAL,
+  - feasiblePoints:
+    - leaky_relu
+    - relu
+    - elu
+    - tanh
+    parameterName: act
+    type: CATEGORICAL
+  name: gcn
 trainer:
-  - parameterName: max_epoch
-    type: INTEGER
-    maxValue: 300
+  hp_space:
+  - maxValue: 300
     minValue: 100
+    parameterName: max_epoch
     scalingType: LINEAR
-  
-  - parameterName: early_stopping_round
     type: INTEGER
-    maxValue: 30
+  - maxValue: 30
     minValue: 10
+    parameterName: early_stopping_round
     scalingType: LINEAR
-
-  - parameterName: lr
-    type: DOUBLE
-    maxValue: 0.05
+    type: INTEGER
+  - maxValue: 0.05
     minValue: 0.01
+    parameterName: lr
     scalingType: LOG
-  
-  - parameterName: weight_decay
     type: DOUBLE
-    maxValue: 0.001
+  - maxValue: 0.001
     minValue: 0.0001
+    parameterName: weight_decay
     scalingType: LOG
-
-hpo:
-  name: random
-  max_evals: 10
-
-ensemble:
-  name: ~
+    type: DOUBLE

examples/graph_classification.py

@@ -1,27 +1,85 @@
+"""
+Example of graph classification on given datasets.
+This version use random split to only show the usage of AutoGraphClassifier.
+Refer to `graph_cv.py` for cross validation evaluation of the whole system
+following paper `A Fair Comparison of Graph Neural Networks for Graph Classification`
+"""
+
 import sys
-sys.path.append('../')
+
+sys.path.append("../")
+import random
+import torch
+import numpy as np
 from autogl.datasets import build_dataset_from_name, utils
 from autogl.solver import AutoGraphClassifier
-from autogl.module import Acc, BaseModel
+from autogl.module import Acc
+from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+
+if __name__ == "__main__":
+    parser = ArgumentParser(
+        "auto graph classification", formatter_class=ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "--dataset",
+        default="mutag",
+        type=str,
+        help="graph classification dataset",
+        choices=["mutag", "imdb-b", "imdb-m", "proteins", "collab"],
+    )
+    parser.add_argument(
+        "--configs", default="../configs/graph_classification.yaml", help="config files"
+    )
+    parser.add_argument("--device", type=int, default=0, help="device to run on")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+
+    args = parser.parse_args()
+    if torch.cuda.is_available():
+        torch.cuda.set_device(args.device)
+    seed = args.seed
+    # set random seed
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    dataset = build_dataset_from_name(args.dataset)
+    if args.dataset.startswith("imdb"):
+        from autogl.module.feature.generators import PYGOneHotDegree
+
+        # get max degree
+        from torch_geometric.utils import degree
+
+        max_degree = 0
+        for data in dataset:
+            deg_max = int(degree(data.edge_index[0], data.num_nodes).max().item())
+            max_degree = max(max_degree, deg_max)
+        dataset = PYGOneHotDegree(max_degree).fit_transform(dataset, inplace=False)
+    elif args.dataset == "collab":
+        from autogl.module.feature.auto_feature import Onlyconst
 
-dataset = build_dataset_from_name('mutag')
-utils.graph_random_splits(dataset, train_ratio=0.4, val_ratio=0.4)
+        dataset = Onlyconst().fit_transform(dataset, inplace=False)
+    utils.graph_random_splits(dataset, train_ratio=0.8, val_ratio=0.1, seed=args.seed)
 
-autoClassifier = AutoGraphClassifier.from_config('../configs/graph_classification.yaml')
+    autoClassifier = AutoGraphClassifier.from_config(args.configs)
 
-# train
-autoClassifier.fit(
-    dataset, 
-    time_limit=3600, 
-    train_split=0.8, 
-    val_split=0.1, 
-    cross_validation=True,
-    cv_split=10, 
-)
-autoClassifier.get_leaderboard().show()
+    # train
+    autoClassifier.fit(dataset, evaluation_method=[Acc], seed=args.seed)
+    autoClassifier.get_leaderboard().show()
 
-print('best single model:\n', autoClassifier.get_leaderboard().get_best_model(0))
+    print("best single model:\n", autoClassifier.get_leaderboard().get_best_model(0))
 
-# test
-predict_result = autoClassifier.predict_proba()
-print(Acc.evaluate(predict_result, dataset.data.y[dataset.test_index].cpu().detach().numpy()))
+    # test
+    predict_result = autoClassifier.predict_proba()
+    print(
+        "test acc %.4f"
+        % (
+            Acc.evaluate(
+                predict_result,
+                dataset.data.y[dataset.test_index].cpu().detach().numpy(),
+            )
+        )
+    )

examples/graph_cv.py

@@ -0,0 +1,96 @@
+"""
+Auto graph classification using cross validation methods proposed in
+paper `A Fair Comparison of Graph Neural Networks for Graph Classification`
+"""
+
+import sys
+import random
+import torch
+import numpy as np
+from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+
+sys.path.append("../")
+
+from autogl.datasets import build_dataset_from_name, utils
+from autogl.solver import AutoGraphClassifier
+from autogl.module import Acc
+
+if __name__ == "__main__":
+    parser = ArgumentParser(
+        "auto graph classification", formatter_class=ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "--dataset",
+        default="mutag",
+        type=str,
+        help="graph classification dataset",
+        choices=["mutag", "imdb-b", "imdb-m", "proteins", "collab"],
+    )
+    parser.add_argument(
+        "--configs", default="../configs/graph_classification.yaml", help="config files"
+    )
+    parser.add_argument("--device", type=int, default=0, help="device to run on")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--folds", type=int, default=10, help="fold number")
+
+    args = parser.parse_args()
+    if torch.cuda.is_available():
+        torch.cuda.set_device(args.device)
+    seed = args.seed
+    # set random seed
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    print("begin processing dataset", args.dataset, "into", args.folds, "folds.")
+    dataset = build_dataset_from_name(args.dataset)
+    if args.dataset.startswith("imdb"):
+        from autogl.module.feature.generators import PYGOneHotDegree
+
+        # get max degree
+        from torch_geometric.utils import degree
+
+        max_degree = 0
+        for data in dataset:
+            deg_max = int(degree(data.edge_index[0], data.num_nodes).max().item())
+            max_degree = max(max_degree, deg_max)
+        dataset = PYGOneHotDegree(max_degree).fit_transform(dataset, inplace=False)
+    elif args.dataset == "collab":
+        from autogl.module.feature.auto_feature import Onlyconst
+
+        dataset = Onlyconst().fit_transform(dataset, inplace=False)
+    utils.graph_cross_validation(dataset, args.folds, random_seed=args.seed)
+
+    accs = []
+    for fold in range(args.folds):
+        print("evaluating on fold number:", fold)
+        utils.graph_set_fold_id(dataset, fold)
+        train_dataset = utils.graph_get_split(dataset, "train", False)
+        autoClassifier = AutoGraphClassifier.from_config(args.configs)
+
+        autoClassifier.fit(
+            train_dataset,
+            train_split=0.9,
+            val_split=0.1,
+            seed=args.seed,
+            evaluation_method=[Acc],
+        )
+        predict_result = autoClassifier.predict_proba(dataset, mask="val")
+        acc = Acc.evaluate(
+            predict_result, dataset.data.y[dataset.val_index].cpu().detach().numpy()
+        )
+        print(
+            "test acc %.4f"
+            % (
+                Acc.evaluate(
+                    predict_result,
+                    dataset.data.y[dataset.val_index].cpu().detach().numpy(),
+                )
+            )
+        )
+        accs.append(acc)
+    print("Average acc on", args.dataset, ":", np.mean(accs), "~", np.std(accs))

examples/node_classification.py

@@ -1,5 +1,6 @@
 import sys
-sys.path.append('../')
+
+sys.path.append("../")
 from autogl.datasets import build_dataset_from_name
 from autogl.solver import AutoNodeClassifier
 from autogl.module import Acc
@@ -8,20 +9,42 @@ import random
 import torch
 import numpy as np
 
-import logging
-logging.basicConfig(level=logging.INFO)
+if __name__ == "__main__":
 
-if __name__ == '__main__':
+    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
 
-    from argparse import ArgumentParser
-    parser = ArgumentParser()
-    parser.add_argument('--dataset', default='cora', type=str)
-    parser.add_argument('--configs', type=str, default='../configs/nodeclf_gcn_benchmark_small.yml')
+    parser = ArgumentParser(
+        "auto node classification", formatter_class=ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "--dataset",
+        default="cora",
+        type=str,
+        help="dataset to use",
+        choices=[
+            "cora",
+            "pubmed",
+            "citeseer",
+            "coauthor_cs",
+            "coauthor_physics",
+            "amazon_computers",
+            "amazon_photo",
+        ],
+    )
+    parser.add_argument(
+        "--configs",
+        type=str,
+        default="../configs/nodeclf_gcn_benchmark_small.yml",
+        help="config to use",
+    )
     # following arguments will override parameters in the config file
-    parser.add_argument('--hpo', type=str, default='random')
-    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)
+    parser.add_argument("--hpo", type=str, default="tpe", help="hpo methods")
+    parser.add_argument(
+        "--max_eval", type=int, default=50, help="max hpo evaluation times"
+    )
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--device", default=0, type=int, help="GPU device")
+
     args = parser.parse_args()
     if torch.cuda.is_available():
         torch.cuda.set_device(args.device)
@@ -36,23 +59,30 @@ if __name__ == '__main__':
         torch.backends.cudnn.benchmark = False
 
     dataset = build_dataset_from_name(args.dataset)
-    
-    configs = yaml.load(open(args.configs, 'r').read(), Loader=yaml.FullLoader)
-    configs['hpo']['name'] = args.hpo
-    configs['hpo']['max_evals'] = args.max_eval
+
+    configs = yaml.load(open(args.configs, "r").read(), Loader=yaml.FullLoader)
+    configs["hpo"]["name"] = args.hpo
+    configs["hpo"]["max_evals"] = args.max_eval
     autoClassifier = AutoNodeClassifier.from_config(configs)
 
     # train
-    if args.dataset in ['cora', 'citeseer', 'pubmed']:
+    if args.dataset in ["cora", "citeseer", "pubmed"]:
         autoClassifier.fit(dataset, time_limit=3600, evaluation_method=[Acc])
     else:
-        autoClassifier.fit(dataset, time_limit=3600, evaluation_method=[Acc], seed=seed, train_split=20*dataset.num_classes, val_split=30*dataset.num_classes, balanced=False)
-    val = autoClassifier.get_model_by_performance(0)[0].get_valid_score()[0]
-    print('val acc: ', val)
+        autoClassifier.fit(
+            dataset,
+            time_limit=3600,
+            evaluation_method=[Acc],
+            seed=seed,
+            train_split=20 * dataset.num_classes,
+            val_split=30 * dataset.num_classes,
+            balanced=False,
+        )
+    autoClassifier.get_leaderboard().show()
 
     # test
-    predict_result = autoClassifier.predict_proba(use_best=True, use_ensemble=False)
-    print('test acc: ', Acc.evaluate(predict_result, dataset.data.y[dataset.data.test_mask].numpy()))
-
-
-
+    predict_result = autoClassifier.predict_proba()
+    print(
+        "test acc: %.4f"
+        % (Acc.evaluate(predict_result, dataset.data.y[dataset.data.test_mask].numpy()))
+    )