diff --git a/autogl/datasets/utils.py b/autogl/datasets/utils.py index 7b5679c..4885ea0 100644 --- a/autogl/datasets/utils.py +++ b/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 = {} diff --git a/autogl/module/model/__init__.py b/autogl/module/model/__init__.py index 42bdbc4..4bff9bb 100644 --- a/autogl/module/model/__init__.py +++ b/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 diff --git a/autogl/module/model/gcn.py b/autogl/module/model/gcn.py index 7f279a7..e28bd80 100644 --- a/autogl/module/model/gcn.py +++ b/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""" diff --git a/autogl/module/train/__init__.py b/autogl/module/train/__init__.py index 8e70ff2..a7fe443 100644 --- a/autogl/module/train/__init__.py +++ b/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", ] diff --git a/autogl/module/train/base.py b/autogl/module/train/base.py index 76c6a61..ba36e74 100644 --- a/autogl/module/train/base.py +++ b/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 - ) \ No newline at end of file + ) diff --git a/autogl/module/train/evaluation.py b/autogl/module/train/evaluation.py index d36f708..db16fd3 100644 --- a/autogl/module/train/evaluation.py +++ b/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 diff --git a/autogl/module/train/link_prediction.py b/autogl/module/train/link_prediction.py index a444ae3..56cf3fe 100644 --- a/autogl/module/train/link_prediction.py +++ b/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 \ No newline at end of file + link_labels[: pos_edge_index.size(1)] = 1.0 + return link_labels diff --git a/autogl/solver/classifier/link_predictor.py b/autogl/solver/classifier/link_predictor.py index 040da36..71fe591 100644 --- a/autogl/solver/classifier/link_predictor.py +++ b/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(