black format

autogl/datasets/utils.py

@@ -9,9 +9,12 @@ from sklearn.model_selection import StratifiedKFold, KFold
 def split_edges(dataset, train_ratio, val_ratio):
     datas = [data for data in dataset]
     for i in range(len(datas)):
-        datas[i] = train_test_split_edges(datas[i], val_ratio, 1 - train_ratio - val_ratio)
+        datas[i] = train_test_split_edges(
+            datas[i], val_ratio, 1 - train_ratio - val_ratio
+        )
     dataset.data, dataset.slices = dataset.collate(datas)
 
+
 def get_label_number(dataset):
     r"""Get the number of labels in this dataset as dict."""
     label_num = {}

autogl/module/model/__init__.py

@@ -1,7 +1,8 @@
 from ._model_registry import MODEL_DICT, ModelUniversalRegistry, register_model
 from .base import BaseModel
 from .topkpool import AutoTopkpool
-#from .graph_sage import AutoSAGE
+
+# from .graph_sage import AutoSAGE
 from .graphsage import AutoSAGE
 from .gcn import AutoGCN
 from .gat import AutoGAT

autogl/module/model/gcn.py

@@ -23,18 +23,12 @@ class GCN(torch.nn.Module):
         self.__convolution_layers: torch.nn.ModuleList = torch.nn.ModuleList()
         num_layers: int = len(hidden_features) + 1
         if num_layers == 1:
-            self.__convolution_layers.append(
-                GCNConv(num_features, num_classes)
-            )
+            self.__convolution_layers.append(GCNConv(num_features, num_classes))
         else:
-            self.__convolution_layers.append(
-                GCNConv(num_features, hidden_features[0])
-            )
+            self.__convolution_layers.append(GCNConv(num_features, hidden_features[0]))
             for i in range(len(hidden_features)):
                 self.__convolution_layers.append(
-                    GCNConv(
-                        hidden_features[i], hidden_features[i + 1]
-                    )
+                    GCNConv(hidden_features[i], hidden_features[i + 1])
                     if i + 1 < len(hidden_features)
                     else GCNConv(hidden_features[i], num_classes)
                 )
@@ -109,9 +103,8 @@ class GCN(torch.nn.Module):
         return (prob_adj > 0).nonzero(as_tuple=False).t()
 
 
-
-#@register_model("gcn")
-#class AutoGCN(ClassificationModel):
+# @register_model("gcn")
+# class AutoGCN(ClassificationModel):
 @register_model("gcn")
 class AutoGCN(BaseModel):
     r"""

autogl/module/train/__init__.py

@@ -4,9 +4,10 @@ from .base import (
     Evaluation,
     BaseNodeClassificationTrainer,
     BaseGraphClassificationTrainer,
-    BaseLinkPredictionTrainer
+    BaseLinkPredictionTrainer,
 )
 
+
 def register_trainer(name):
     def register_trainer_cls(cls):
         if name in TRAINER_DICT:
@@ -20,6 +21,7 @@ def register_trainer(name):
 
     return register_trainer_cls
 
+
 from .graph_classification_full import GraphClassificationFullTrainer
 from .node_classification_full import NodeClassificationFullTrainer
 from .link_prediction import LinkPredictionTrainer
@@ -39,5 +41,5 @@ __all__ = [
     "Auc",
     "Logloss",
     "Mrr",
-    "get_feval"
+    "get_feval",
 ]

autogl/module/train/base.py

@@ -403,6 +403,7 @@ class BaseGraphClassificationTrainer(_BaseClassificationTrainer):
             model, num_features, num_classes, device, init, feval, loss
         )
 
+
 class BaseLinkPredictionTrainer(_BaseClassificationTrainer):
     def __init__(
         self,
@@ -417,4 +418,4 @@ class BaseLinkPredictionTrainer(_BaseClassificationTrainer):
     ):
         super(BaseLinkPredictionTrainer, self).__init__(
             model, num_features, 2, device, init, feval, loss
-        )
+        )

autogl/module/train/evaluation.py

@@ -120,7 +120,9 @@ class Auc(Evaluation):
         if len(predict.shape) == 1:
             pos_predict = predict
         else:
-            assert predict.shape[1] == 2, "Cannot use auc on given data with %d classes!" % (predict.shape[1])
+            assert (
+                predict.shape[1] == 2
+            ), "Cannot use auc on given data with %d classes!" % (predict.shape[1])
             pos_predict = predict[:, 1]
         return roc_auc_score(label, pos_predict)
 
@@ -169,7 +171,9 @@ class Mrr(Evaluation):
         Should return: the evaluation result (float)
         """
         if len(predict.shape) == 2:
-            assert predict.shape[1] == 2, "Cannot use mrr on given data with %d classes!" % (predict.shape[1])
+            assert (
+                predict.shape[1] == 2
+            ), "Cannot use mrr on given data with %d classes!" % (predict.shape[1])
             pos_predict = predict[:, 1]
         else:
             pos_predict = predict

autogl/module/train/link_prediction.py

@@ -13,6 +13,7 @@ from ...utils import get_logger
 
 LOGGER = get_logger("link prediction trainer")
 
+
 def get_feval(feval):
     if isinstance(feval, str):
         return EVALUATE_DICT[feval]
@@ -65,12 +66,12 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
         max_epoch=100,
         early_stopping_round=101,
         weight_decay=1e-4,
-        device='auto',
+        device="auto",
         init=True,
         feval=[Auc],
         loss="binary_cross_entropy_with_logits",
         *args,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(model, num_features, device, init, feval, loss)
 
@@ -189,19 +190,27 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
             self.model.model.train()
 
             neg_edge_index = negative_sampling(
-                edge_index=data.train_pos_edge_index, num_nodes=data.num_nodes,
-                num_neg_samples=data.train_pos_edge_index.size(1))
+                edge_index=data.train_pos_edge_index,
+                num_nodes=data.num_nodes,
+                num_neg_samples=data.train_pos_edge_index.size(1),
+            )
 
             optimizer.zero_grad()
             # res = self.model.model.forward(data)
             z = self.model.model.encode(data)
-            link_logits = self.model.model.decode(z, data.train_pos_edge_index, neg_edge_index)
-            link_labels = self.get_link_labels(data.train_pos_edge_index, neg_edge_index)
+            link_logits = self.model.model.decode(
+                z, data.train_pos_edge_index, neg_edge_index
+            )
+            link_labels = self.get_link_labels(
+                data.train_pos_edge_index, neg_edge_index
+            )
             # loss = F.binary_cross_entropy_with_logits(link_logits, link_labels)
             if hasattr(F, self.loss):
                 loss = getattr(F, self.loss)(link_logits, link_labels)
             else:
-                raise TypeError("PyTorch does not support loss type {}".format(self.loss))
+                raise TypeError(
+                    "PyTorch does not support loss type {}".format(self.loss)
+                )
 
             loss.backward()
             optimizer.step()
@@ -211,7 +220,7 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
                 feval = self.feval[0]
             else:
                 feval = self.feval
-            val_loss = self.evaluate([data], mask='val', feval=feval)
+            val_loss = self.evaluate([data], mask="val", feval=feval)
             if feval.is_higher_better() is True:
                 val_loss = -val_loss
             self.early_stopping(val_loss, self.model.model)
@@ -261,10 +270,8 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
         self.train_only(data)
         if keep_valid_result:
             self.valid_result = self.predict_only(data)
-            self.valid_result_prob = self.predict_proba(dataset, 'val')
-            self.valid_score = self.evaluate(
-                dataset, mask='val', feval=self.feval
-            )
+            self.valid_result_prob = self.predict_proba(dataset, "val")
+            self.valid_score = self.evaluate(dataset, mask="val", feval=self.feval)
 
     def predict(self, dataset, mask=None):
         """
@@ -304,11 +311,11 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
         data = dataset[0]
         data = data.to(self.device)
         if mask in ["train", "val", "test"]:
-            pos_edge_index = data[f'{mask}_pos_edge_index']
-            neg_edge_index = data[f'{mask}_neg_edge_index']
+            pos_edge_index = data[f"{mask}_pos_edge_index"]
+            neg_edge_index = data[f"{mask}_neg_edge_index"]
         else:
-            pos_edge_index = data[f'test_pos_edge_index']
-            neg_edge_index = data[f'test_neg_edge_index']
+            pos_edge_index = data[f"test_pos_edge_index"]
+            neg_edge_index = data[f"test_neg_edge_index"]
 
         self.model.model.eval()
         with torch.no_grad():
@@ -400,11 +407,11 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
             feval = get_feval(feval)
 
         if mask in ["train", "val", "test"]:
-            pos_edge_index = data[f'{mask}_pos_edge_index']
-            neg_edge_index = data[f'{mask}_neg_edge_index']
+            pos_edge_index = data[f"{mask}_pos_edge_index"]
+            neg_edge_index = data[f"{mask}_neg_edge_index"]
         else:
-            pos_edge_index = data[f'test_pos_edge_index']
-            neg_edge_index = data[f'test_neg_edge_index']
+            pos_edge_index = data[f"test_pos_edge_index"]
+            neg_edge_index = data[f"test_neg_edge_index"]
 
         self.model.model.eval()
         with torch.no_grad():
@@ -480,7 +487,7 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
             feval=self.feval,
             init=True,
             *self.args,
-            **self.kwargs
+            **self.kwargs,
         )
 
         return ret
@@ -507,5 +514,5 @@ class LinkPredictionTrainer(BaseLinkPredictionTrainer):
     def get_link_labels(self, pos_edge_index, neg_edge_index):
         E = pos_edge_index.size(1) + neg_edge_index.size(1)
         link_labels = torch.zeros(E, dtype=torch.float, device=self.device)
-        link_labels[:pos_edge_index.size(1)] = 1.
-        return link_labels
+        link_labels[: pos_edge_index.size(1)] = 1.0
+        return link_labels

autogl/solver/classifier/link_predictor.py

@@ -195,8 +195,8 @@ class AutoLinkPredictor(BaseClassifier):
     def _to_prob(self, sig_prob: np.ndarray):
         nelements = len(sig_prob)
         prob = np.zeros([nelements, 2])
-        prob[:,0] = 1 - sig_prob
-        prob[:,1] = sig_prob
+        prob[:, 0] = 1 - sig_prob
+        prob[:, 1] = sig_prob
         return prob
 
     # pylint: disable=arguments-differ
@@ -277,14 +277,19 @@ class AutoLinkPredictor(BaseClassifier):
         if train_split is not None and val_split is not None:
             utils.split_edges(dataset, train_split, val_split)
         else:
-            assert all([hasattr(dataset.data, f'{name}') for name in [
-                'train_pos_edge_index', 
-                'train_neg_adj_mask', 
-                'val_pos_edge_index',
-                'val_neg_edge_index', 
-                'test_pos_edge_index', 
-                'test_neg_edge_index'
-            ]]), (
+            assert all(
+                [
+                    hasattr(dataset.data, f"{name}")
+                    for name in [
+                        "train_pos_edge_index",
+                        "train_neg_adj_mask",
+                        "val_pos_edge_index",
+                        "val_neg_edge_index",
+                        "test_pos_edge_index",
+                        "test_neg_edge_index",
+                    ]
+                ]
+            ), (
                 "The dataset has no default train/val split! Please manually pass "
                 "train and val ratio."
             )
@@ -307,7 +312,9 @@ class AutoLinkPredictor(BaseClassifier):
             num_features=self.dataset[0].x.shape[1],
             feval=evaluator_list,
             device=self.runtime_device,
-            loss="binary_cross_entropy_with_logits" if not hasattr(dataset, "loss") else dataset.loss,
+            loss="binary_cross_entropy_with_logits"
+            if not hasattr(dataset, "loss")
+            else dataset.loss,
         )
 
         # train the models and tune hpo
@@ -330,7 +337,9 @@ class AutoLinkPredictor(BaseClassifier):
             name = optimized.get_name_with_hp() + "_idx%d" % (idx)
             names.append(name)
             performance_on_valid, _ = optimized.get_valid_score(return_major=False)
-            result_valid.append(self._to_prob(optimized.get_valid_predict_proba().cpu().numpy()))
+            result_valid.append(
+                self._to_prob(optimized.get_valid_predict_proba().cpu().numpy())
+            )
             self.leaderboard.insert_model_performance(
                 name,
                 dict(
@@ -344,10 +353,13 @@ class AutoLinkPredictor(BaseClassifier):
 
         # fit the ensemble model
         if self.ensemble_module is not None:
-            pos_edge_index, neg_edge_index = self.dataset[0].val_pos_edge_index, self.dataset[0].val_neg_edge_index
+            pos_edge_index, neg_edge_index = (
+                self.dataset[0].val_pos_edge_index,
+                self.dataset[0].val_neg_edge_index,
+            )
             E = pos_edge_index.size(1) + neg_edge_index.size(1)
             link_labels = torch.zeros(E, dtype=torch.float)
-            link_labels[:pos_edge_index.size(1)] = 1.
+            link_labels[: pos_edge_index.size(1)] = 1.0
 
             performance = self.ensemble_module.fit(
                 result_valid,
@@ -519,10 +531,12 @@ class AutoLinkPredictor(BaseClassifier):
             names = []
             for model_name in self.trained_models:
                 predict_result.append(
-                    self._to_prob(self._predict_proba_by_name(dataset, model_name, mask))
+                    self._to_prob(
+                        self._predict_proba_by_name(dataset, model_name, mask)
+                    )
                 )
                 names.append(model_name)
-            return self.ensemble_module.ensemble(predict_result, names)[:,1]
+            return self.ensemble_module.ensemble(predict_result, names)[:, 1]
 
         if use_ensemble and self.ensemble_module is None:
             LOGGER.warning(