diff --git a/autogl/datasets/_heterogeneous_datasets/_dgl_heterogeneous_datasets.py b/autogl/datasets/_heterogeneous_datasets/_dgl_heterogeneous_datasets.py index 4c314a1..133f75b 100644 --- a/autogl/datasets/_heterogeneous_datasets/_dgl_heterogeneous_datasets.py +++ b/autogl/datasets/_heterogeneous_datasets/_dgl_heterogeneous_datasets.py @@ -10,7 +10,7 @@ from .. import _dataset_registry def get_binary_mask(total_size, indices): mask = torch.zeros(total_size) mask[indices] = 1 - return mask.byte() + return mask.bool() @_dataset_registry.DatasetUniversalRegistry.register_dataset("hetero-acm-han") diff --git a/autogl/module/model/dgl/base.py b/autogl/module/model/dgl/base.py index 5d0840b..6ad7ede 100644 --- a/autogl/module/model/dgl/base.py +++ b/autogl/module/model/dgl/base.py @@ -35,6 +35,9 @@ class BaseAutoModel(AutoModule): self._kwargs = kwargs super(BaseAutoModel, self).__init__(device) + def to(self, device): + return self.to_device(device) + def to_device(self, device): self.device = device if self.model is not None: diff --git a/autogl/module/model/dgl/hetero/han.py b/autogl/module/model/dgl/hetero/han.py index a31ff0a..f8796ec 100644 --- a/autogl/module/model/dgl/hetero/han.py +++ b/autogl/module/model/dgl/hetero/han.py @@ -133,7 +133,7 @@ class HAN(nn.Module): self.layers.append(HANLayer(self.args["meta_paths"], self.args["num_features"], self.args["hidden"][0], self.args["heads"][0], self.args["dropout"], act)) for l in range(1, len(self.args["heads"])): self.layers.append(HANLayer(self.args["meta_paths"], self.args["hidden"][l-1] * self.args["heads"][l-1], - self.args["hidden"], self.args["heads"][l], self.args["dropout"], act)) + self.args["hidden"][l], self.args["heads"][l], self.args["dropout"], act)) self.predict = nn.Linear(self.args["hidden"][-1] * self.args["heads"][-1], self.args["num_class"]) def forward(self, g): @@ -206,7 +206,10 @@ class AutoHAN(BaseHeteroModelMaintainer): "parameterName": "heads", "type": "NUMERICAL_LIST", "numericalType": "INTEGER", - "feasiblePoints": "[8]", + "scalingType": "LOG", + "length": 3, + "minValue": [1, 1, 1], + "maxValue": [16, 16, 16], "cutPara": ("num_layers",), "cutFunc": lambda x: x[0] - 1, }, diff --git a/autogl/module/train/__init__.py b/autogl/module/train/__init__.py index cf428b3..c83f1c3 100644 --- a/autogl/module/train/__init__.py +++ b/autogl/module/train/__init__.py @@ -7,6 +7,7 @@ from .base import ( BaseNodeClassificationTrainer, BaseGraphClassificationTrainer, BaseLinkPredictionTrainer, + BaseNodeClassificationHetTrainer ) @@ -36,6 +37,7 @@ __all__ = [ "Evaluation", "BaseGraphClassificationTrainer", "BaseNodeClassificationTrainer", + "BaseNodeClassificationHetTrainer", "BaseLinkPredictionTrainer", "GraphClassificationFullTrainer", "NodeClassificationFullTrainer", diff --git a/autogl/module/train/node_classification_het.py b/autogl/module/train/node_classification_het.py index aaf4f78..6c5b035 100644 --- a/autogl/module/train/node_classification_het.py +++ b/autogl/module/train/node_classification_het.py @@ -71,7 +71,7 @@ class NodeClassificationHetTrainer(BaseNodeClassificationHetTrainer): early_stopping_round=100, weight_decay=1e-4, device="auto", - init=True, + init=False, feval=[Logloss], loss="nll_loss", lr_scheduler_type=None, @@ -343,7 +343,7 @@ class NodeClassificationHetTrainer(BaseNodeClassificationHetTrainer): def _get_mask(self, dataset, mask): if mask in ["train", "val", "test"]: - return dataset[0].nodes[dataset.schema["target_node_type"]].data[f"{mask}_mask"].bool() + return dataset[0].nodes[dataset.schema["target_node_type"]].data[f"{mask}_mask"] return mask def evaluate(self, dataset, mask='val', feval = None): diff --git a/autogl/solver/__init__.py b/autogl/solver/__init__.py index 1cba0b1..5896f5e 100644 --- a/autogl/solver/__init__.py +++ b/autogl/solver/__init__.py @@ -2,12 +2,13 @@ Auto solver for various graph tasks """ -from .classifier import AutoGraphClassifier, AutoNodeClassifier, AutoLinkPredictor +from .classifier import AutoGraphClassifier, AutoNodeClassifier, AutoLinkPredictor, AutoHeteroNodeClassifier from .utils import LeaderBoard __all__ = [ "AutoNodeClassifier", "AutoGraphClassifier", "AutoLinkPredictor", + "AutoHeteroNodeClassifier", "LeaderBoard", ] diff --git a/autogl/solver/classifier/__init__.py b/autogl/solver/classifier/__init__.py index e30c582..c82dea5 100644 --- a/autogl/solver/classifier/__init__.py +++ b/autogl/solver/classifier/__init__.py @@ -6,10 +6,16 @@ from .base import BaseClassifier from .graph_classifier import AutoGraphClassifier from .node_classifier import AutoNodeClassifier from .link_predictor import AutoLinkPredictor +from autogl.backend import DependentBackend +if DependentBackend.is_dgl(): + from .hetero import AutoHeteroNodeClassifier +else: + AutoHeteroNodeClassifier = None __all__ = [ "BaseClassifier", "AutoGraphClassifier", "AutoNodeClassifier", "AutoLinkPredictor", + "AutoHeteroNodeClassifier" ] diff --git a/autogl/solver/classifier/hetero/__init__.py b/autogl/solver/classifier/hetero/__init__.py new file mode 100644 index 0000000..fbffed0 --- /dev/null +++ b/autogl/solver/classifier/hetero/__init__.py @@ -0,0 +1 @@ +from .node_classifier import AutoHeteroNodeClassifier \ No newline at end of file diff --git a/autogl/solver/classifier/hetero/node_classifier.py b/autogl/solver/classifier/hetero/node_classifier.py new file mode 100644 index 0000000..9f944b2 --- /dev/null +++ b/autogl/solver/classifier/hetero/node_classifier.py @@ -0,0 +1,669 @@ +""" +Auto Classfier for Heterogeneous Node Classification +""" +import time +import json + +from copy import deepcopy +from typing import Sequence + +import torch +import numpy as np +import yaml + +from ..base import BaseClassifier +from ...base import _parse_hp_space, _initialize_single_model, _parse_model_hp +from ....module.train import TRAINER_DICT, BaseNodeClassificationHetTrainer +from ....module.train import get_feval +from ...utils import LeaderBoard, set_seed +from ....utils import get_logger + +LOGGER = get_logger("HeteroNodeClassifier") + +class AutoHeteroNodeClassifier(BaseClassifier): + """ + Auto Multi-class HeteroGraph Node Classifier. + + Used to automatically solve the heterogeneous node classification problems. + + Parameters + ---------- + feature_module: autogl.module.feature.BaseFeatureEngineer or str or None + The (name of) auto feature engineer used to process the given dataset. Default ``deepgl``. + Disable feature engineer by setting it to ``None``. + + graph_models: list of autogl.module.model.BaseModel or list of str + The (name of) models to be optimized as backbone. Default ``['gat', 'gcn']``. + + hpo_module: autogl.module.hpo.BaseHPOptimizer or str or None + The (name of) hpo module used to search for best hyper parameters. Default ``anneal``. + Disable hpo by setting it to ``None``. + + ensemble_module: autogl.module.ensemble.BaseEnsembler or str or None + The (name of) ensemble module used to ensemble the multi-models found. Default ``voting``. + Disable ensemble by setting it to ``None``. + + max_evals: int (Optional) + If given, will set the number eval times the hpo module will use. + Only be effective when hpo_module is ``str``. 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. + If a list of trainer hp is given, will specify every model with corrsponding + trainer hp space. + Default ``None``. + + model_hp_spaces: list of list of dict (Optional) + model hp space configuration. + If given, will specify every hp space of every passed model. Default ``None``. + + size: int (Optional) + The max models ensemble module will use. Default ``None``. + + device: torch.device or str + The device where model will be running on. If set to ``auto``, will use gpu when available. + You can also specify the device by directly giving ``gpu`` or ``cuda:0``, etc. + Default ``auto``. + """ + + def __init__( + self, + graph_models=("han", "hgt"), + hpo_module="anneal", + ensemble_module="voting", + max_evals=50, + default_trainer="NodeClassificationHet", + trainer_hp_space=None, + model_hp_spaces=None, + size=4, + device="auto", + ): + + super().__init__( + # currently we do not support feature engineering + feature_module=None, + graph_models=graph_models, + # currently we do not support nas for heterogeneous node classifier + nas_algorithms=None, + nas_spaces=None, + nas_estimators=None, + hpo_module=hpo_module, + ensemble_module=ensemble_module, + max_evals=max_evals, + default_trainer=default_trainer, + trainer_hp_space=trainer_hp_space, + model_hp_spaces=model_hp_spaces, + size=size, + device=device, + ) + + # data to be kept when fit + self.dataset = None + + def _init_graph_module( + self, graph_models, num_classes, num_features, feval, device, loss, dataset + ) -> "AutoHeteroNodeClassifier": + # load graph network module + self.graph_model_list = [] + + for i, model in enumerate(graph_models): + # init the trainer + if not isinstance(model, BaseNodeClassificationHetTrainer): + trainer = ( + self._default_trainer if not isinstance(self._default_trainer, (tuple, list)) + else self._default_trainer[i] + ) + if isinstance(trainer, str): + trainer = TRAINER_DICT[trainer]() + if isinstance(model, (tuple, list)): + trainer.encoder = model[0] + trainer.decoder = model[1] + else: + trainer.encoder = model + else: + trainer = model + + # set model hp space + if self._model_hp_spaces is not None: + if self._model_hp_spaces[i] is not None: + if isinstance(self._model_hp_spaces[i], dict): + encoder_hp_space = self._model_hp_spaces[i].get('encoder', None) + decoder_hp_space = self._model_hp_spaces[i].get('decoder', None) + else: + encoder_hp_space = self._model_hp_spaces[i] + decoder_hp_space = None + if encoder_hp_space is not None: + trainer.encoder.hyper_parameter_space = encoder_hp_space + if decoder_hp_space is not None: + trainer.decoder.hyper_parameter_space = decoder_hp_space + + # 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 + + trainer.num_features = num_features + trainer.num_classes = num_classes + trainer.from_dataset(dataset) + trainer.loss = loss + trainer.feval = feval + trainer.to(device) + self.graph_model_list.append(trainer) + + return self + + # pylint: disable=arguments-differ + def fit( + self, + dataset, + time_limit=-1, + evaluation_method="infer", + seed=None, + ) -> "AutoHeteroNodeClassifier": + """ + Fit current solver on given dataset. + + Parameters + ---------- + dataset: autogl.data.Dataset + The dataset needed to fit on. This dataset must have only one graph. + + time_limit: int + The time limit of the whole fit process (in seconds). If set below 0, + will ignore time limit. Default ``-1``. + + inplace: bool + Whether we process the given dataset in inplace manner. Default ``False``. + Set it to True if you want to save memory by modifying the given dataset directly. + + evaluation_method: (list of) str or autogl.module.train.evaluation + A (list of) evaluation method for current solver. If ``infer``, will automatically + determine. Default ``infer``. + + seed: int (Optional) + The random seed. If set to ``None``, will run everything at random. + Default ``None``. + + Returns + ------- + self: autogl.solver.AutoNodeClassifier + A reference of current solver. + """ + set_seed(seed) + + if time_limit < 0: + time_limit = 3600 * 24 + time_begin = time.time() + + graph_data = dataset[0] + field = dataset.schema["target_node_type"] + all_labels = graph_data.nodes[field].data['label'] + num_classes = all_labels.max().item() + 1 + + # initialize leaderboard + if evaluation_method == "infer": + if hasattr(dataset, "metric"): + evaluation_method = [dataset.metric] + else: + num_of_label = num_classes + if num_of_label == 2: + evaluation_method = ["auc"] + else: + evaluation_method = ["acc"] + assert isinstance(evaluation_method, list) + evaluator_list = get_feval(evaluation_method) + + self.leaderboard = LeaderBoard( + [e.get_eval_name() for e in evaluator_list], + {e.get_eval_name(): e.is_higher_better() for e in evaluator_list}, + ) + + # set up the dataset + assert ("train_mask" in graph_data.nodes[field].data + and "val_mask" in graph_data.nodes[field].data), ("Currently only support" + " Dataset with default train/val/test split") + + self.dataset = dataset + + # check whether the dataset has features. + # currently we only support hetero graph classification with features. + + feat = graph_data.nodes[field].data['feat'] + assert feat is not None, ( + "Does not support fit on non node-feature dataset!" + " Please add node features to dataset or specify feature engineers that generate" + " node features." + ) + + num_features = feat.size(-1) + + # initialize graph networks + self._init_graph_module( + self.gml, + num_features=num_features, + num_classes=num_classes, + feval=evaluator_list, + device=self.runtime_device, + loss="nll_loss" if not hasattr(dataset, "loss") else self.dataset.loss, + dataset=dataset + ) + + # train the models and tune hpo + result_valid = [] + names = [] + for idx, model in enumerate(self.graph_model_list): + time_for_each_model = (time_limit - time.time() + time_begin) / ( + len(self.graph_model_list) - idx + ) + if self.hpo_module is None: + model.initialize() + model.train(dataset, True) + optimized = model + else: + optimized, _ = self.hpo_module.optimize( + trainer=model, dataset=dataset, time_limit=time_for_each_model + ) + # to save memory, all the trainer derived will be mapped to cpu + optimized.to(torch.device("cpu")) + name = str(optimized) + "_idx%d" % (idx) + names.append(name) + performance_on_valid, _ = optimized.get_valid_score(return_major=False) + result_valid.append(optimized.get_valid_predict_proba().cpu().numpy()) + self.leaderboard.insert_model_performance( + name, + dict( + zip( + [e.get_eval_name() for e in evaluator_list], + performance_on_valid, + ) + ), + ) + self.trained_models[name] = optimized + + # fit the ensemble model + if self.ensemble_module is not None: + performance = self.ensemble_module.fit( + result_valid, + all_labels[graph_data.nodes[field].data["val_mask"]].cpu().numpy(), + names, + evaluator_list, + n_classes=num_classes, + ) + self.leaderboard.insert_model_performance( + "ensemble", + dict(zip([e.get_eval_name() for e in evaluator_list], performance)), + ) + + return self + + def fit_predict( + self, + dataset, + time_limit=-1, + evaluation_method="infer", + use_ensemble=True, + use_best=True, + name=None, + ) -> np.ndarray: + """ + Fit current solver on given dataset and return the predicted value. + + Parameters + ---------- + dataset: torch_geometric.data.dataset.Dataset + The dataset needed to fit on. This dataset must have only one graph. + + time_limit: int + The time limit of the whole fit process (in seconds). + If set below 0, will ignore time limit. Default ``-1``. + + inplace: bool + Whether we process the given dataset in inplace manner. Default ``False``. + Set it to True if you want to save memory by modifying the given dataset directly. + + train_split: float or int (Optional) + The train ratio (in ``float``) or number (in ``int``) of dataset. If you want to + use default train/val/test split in dataset, please set this to ``None``. + Default ``None``. + + val_split: float or int (Optional) + The validation ratio (in ``float``) or number (in ``int``) of dataset. If you want + to use default train/val/test split in dataset, please set this to ``None``. + Default ``None``. + + balanced: bool + Wether to create the train/valid/test split in a balanced way. + If set to ``True``, the train/valid will have the same number of different classes. + Default ``False``. + + evaluation_method: (list of) str or autogl.module.train.evaluation + A (list of) evaluation method for current solver. If ``infer``, will automatically + determine. Default ``infer``. + + use_ensemble: bool + Whether to use ensemble to do the predict. Default ``True``. + + use_best: bool + Whether to use the best single model to do the predict. Will only be effective when + ``use_ensemble`` is ``False``. + Default ``True``. + + name: str or None + The name of model used to predict. Will only be effective when ``use_ensemble`` and + ``use_best`` both are ``False``. + Default ``None``. + + Returns + ------- + result: np.ndarray + An array of shape ``(N,)``, where ``N`` is the number of test nodes. The prediction + on given dataset. + """ + self.fit( + dataset=dataset, + time_limit=time_limit, + evaluation_method=evaluation_method, + ) + return self.predict( + dataset=dataset, + use_ensemble=use_ensemble, + use_best=use_best, + name=name, + ) + + def predict_proba( + self, + dataset=None, + use_ensemble=True, + use_best=True, + name=None, + mask="test", + ) -> np.ndarray: + """ + Predict the node probability. + + Parameters + ---------- + dataset: torch_geometric.data.dataset.Dataset or None + The dataset needed to predict. If ``None``, will use the processed dataset passed + to ``fit()`` instead. Default ``None``. + + inplaced: bool + Whether the given dataset is processed. Only be effective when ``dataset`` + is not ``None``. If you pass the dataset to ``fit()`` with ``inplace=True``, and + you pass the dataset again to this method, you should set this argument to ``True``. + Otherwise ``False``. Default ``False``. + + inplace: bool + Whether we process the given dataset in inplace manner. Default ``False``. Set it to + True if you want to save memory by modifying the given dataset directly. + + use_ensemble: bool + Whether to use ensemble to do the predict. Default ``True``. + + use_best: bool + Whether to use the best single model to do the predict. Will only be effective when + ``use_ensemble`` is ``False``. Default ``True``. + + name: str or None + The name of model used to predict. Will only be effective when ``use_ensemble`` and + ``use_best`` both are ``False``. Default ``None``. + + mask: str + The data split to give prediction on. Default ``test``. + + Returns + ------- + result: np.ndarray + An array of shape ``(N,C,)``, where ``N`` is the number of test nodes and ``C`` is + the number of classes. The prediction on given dataset. + """ + if dataset is None: + dataset = self.dataset + assert dataset is not None, ( + "Please execute fit() first before" " predicting on remembered dataset" + ) + + if use_ensemble: + LOGGER.info("Ensemble argument on, will try using ensemble model.") + + if not use_ensemble and use_best: + LOGGER.info( + "Ensemble argument off and best argument on, will try using best model." + ) + + if (use_ensemble and self.ensemble_module is not None) or ( + not use_best and name == "ensemble" + ): + # we need to get all the prediction of every model trained + predict_result = [] + names = [] + for model_name in self.trained_models: + predict_result.append( + self._predict_proba_by_name(dataset, model_name, mask) + ) + names.append(model_name) + return self.ensemble_module.ensemble(predict_result, names) + + if use_ensemble and self.ensemble_module is None: + LOGGER.warning( + "Cannot use ensemble because no ensebmle module is given." + " Will use best model instead." + ) + + if use_best or (use_ensemble and self.ensemble_module is None): + # just return the best model we have found + name = self.leaderboard.get_best_model() + return self._predict_proba_by_name(dataset, name, mask) + + if name is not None: + # return model performance by name + return self._predict_proba_by_name(dataset, name, mask) + + LOGGER.error( + "No model name is given while ensemble and best arguments are off." + ) + raise ValueError( + "You need to specify a model name if you do not want use ensemble and best model." + ) + + def _predict_proba_by_name(self, dataset, name, mask="test"): + self.trained_models[name].to(self.runtime_device) + predicted = ( + self.trained_models[name].predict_proba(dataset, mask=mask).cpu().numpy() + ) + self.trained_models[name].to(torch.device("cpu")) + return predicted + + def predict( + self, + dataset=None, + use_ensemble=True, + use_best=True, + name=None, + mask="test", + ) -> np.ndarray: + """ + Predict the node class number. + + Parameters + ---------- + dataset: torch_geometric.data.dataset.Dataset or None + The dataset needed to predict. If ``None``, will use the processed dataset passed + to ``fit()`` instead. Default ``None``. + + inplaced: bool + Whether the given dataset is processed. Only be effective when ``dataset`` + is not ``None``. If you pass the dataset to ``fit()`` with ``inplace=True``, + and you pass the dataset again to this method, you should set this argument + to ``True``. Otherwise ``False``. Default ``False``. + + inplace: bool + Whether we process the given dataset in inplace manner. Default ``False``. + Set it to True if you want to save memory by modifying the given dataset directly. + + use_ensemble: bool + Whether to use ensemble to do the predict. Default ``True``. + + use_best: bool + Whether to use the best single model to do the predict. Will only be effective + when ``use_ensemble`` is ``False``. Default ``True``. + + name: str or None + The name of model used to predict. Will only be effective when ``use_ensemble`` + and ``use_best`` both are ``False``. Default ``None``. + + mask: str + The data split to give prediction on. Default ``test``. + + Returns + ------- + result: np.ndarray + An array of shape ``(N,)``, where ``N`` is the number of test nodes. + The prediction on given dataset. + """ + proba = self.predict_proba( + dataset, use_ensemble, use_best, name, mask + ) + return np.argmax(proba, axis=1) + + + def evaluate(self, dataset=None, + use_ensemble=True, + use_best=True, + name=None, + mask="test", + label=None, + metric="acc" + ): + predicted = self.predict_proba(dataset, use_ensemble, use_best, name, mask) + if dataset is None: + dataset = self.dataset + if label is None: + graph_nodes = dataset[0].nodes[dataset.schema["target_node_type"]].data + if mask in ["train", "val", "test"]: + mask = graph_nodes[f"{mask}_mask"] + label = graph_nodes["label"][mask].cpu().numpy() + evaluator = get_feval(metric) + if isinstance(evaluator, Sequence): + return [evals.evaluate(predicted, label) for evals in evaluator] + return evaluator.evaluate(predicted, label) + + @classmethod + def from_config(cls, path_or_dict, filetype="auto") -> "AutoHeteroNodeClassifier": + """ + Load solver from config file. + + You can use this function to directly load a solver from predefined config dict + or config file path. Currently, only support file type of ``json`` or ``yaml``, + if you pass a path. + + Parameters + ---------- + path_or_dict: str or dict + The path to the config file or the config dictionary object + + filetype: str + The filetype the given file if the path is specified. Currently only support + ``json`` or ``yaml``. You can set to ``auto`` to automatically detect the file + type (from file name). Default ``auto``. + + Returns + ------- + solver: autogl.solver.AutoGraphClassifier + The solver that is created from given file or dictionary. + """ + assert filetype in ["auto", "yaml", "json"], ( + "currently only support yaml file or json file type, but get type " + + filetype + ) + if isinstance(path_or_dict, str): + if filetype == "auto": + if path_or_dict.endswith(".yaml") or path_or_dict.endswith(".yml"): + filetype = "yaml" + elif path_or_dict.endswith(".json"): + filetype = "json" + else: + LOGGER.error( + "cannot parse the type of the given file name, " + "please manually set the file type" + ) + raise ValueError( + "cannot parse the type of the given file name, " + "please manually set the file type" + ) + if filetype == "yaml": + path_or_dict = yaml.load( + open(path_or_dict, "r").read(), Loader=yaml.FullLoader + ) + else: + path_or_dict = json.load(open(path_or_dict, "r")) + + path_or_dict = deepcopy(path_or_dict) + solver = cls(None, [], None, None) + + models = path_or_dict.pop("models", [{"name": "hgt"}, {"name": "han"}]) + # models should be a list of model + # with each element in two cases + # * a dict describing a certain model + # * a dict containing {"encoder": encoder, "decoder": decoder} + model_hp_space = [ + _parse_model_hp(model) for model in models + ] + model_list = [ + _initialize_single_model(model) for model in models + ] + + trainer = path_or_dict.pop("trainer", None) + default_trainer = "NodeClassificationHet" + trainer_space = None + if isinstance(trainer, dict): + # global default + default_trainer = trainer.pop("name", "NodeClassificationHet") + 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", "NodeClassificationHet") + 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: + name = hpo_dict.pop("name") + solver.set_hpo_module(name, **hpo_dict) + + ensemble_dict = path_or_dict.pop("ensemble", {"name": "voting"}) + if ensemble_dict is not None: + name = ensemble_dict.pop("name") + solver.set_ensemble_module(name, **ensemble_dict) + + return solver diff --git a/autogl/solver/classifier/link_predictor.py b/autogl/solver/classifier/link_predictor.py index d6b6c47..e4ca2a0 100644 --- a/autogl/solver/classifier/link_predictor.py +++ b/autogl/solver/classifier/link_predictor.py @@ -12,7 +12,6 @@ import torch import numpy as np import yaml -from ...data import Data from .base import BaseClassifier from ..base import _parse_hp_space, _initialize_single_model, _parse_model_hp from ...module.feature import FEATURE_DICT diff --git a/autogl/solver/classifier/node_classifier.py b/autogl/solver/classifier/node_classifier.py index 5b86e8a..cd04dd6 100644 --- a/autogl/solver/classifier/node_classifier.py +++ b/autogl/solver/classifier/node_classifier.py @@ -5,7 +5,7 @@ import time import json from copy import deepcopy -from typing import Union, Sequence, Tuple +from typing import Sequence import torch import numpy as np @@ -14,7 +14,6 @@ import yaml from .base import BaseClassifier from ..base import _parse_hp_space, _initialize_single_model, _parse_model_hp from ...module.feature import FEATURE_DICT -from ...module.model import BaseEncoderMaintainer, BaseDecoderMaintainer, BaseAutoModel from ...module.train import TRAINER_DICT, BaseNodeClassificationTrainer from ...module.train import get_feval from ...module.nas.space import NAS_SPACE_DICT diff --git a/autogl/solver/utils.py b/autogl/solver/utils.py index cb3430b..58e5ecc 100644 --- a/autogl/solver/utils.py +++ b/autogl/solver/utils.py @@ -269,3 +269,10 @@ def set_seed(seed=None): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False + +def get_graph_labels_hetero(graph, target_node_type): + if isinstance(graph, GeneralStaticGraph): + if 'label' in graph.nodes[target_node_type].data and BACKEND == 'dgl': + return graph.nodes[target_node_type].data['label'] + return None + if BACKEND == 'dgl' and 'label' in graph.ndata[target_node_type]: return graph.ndata[target_node_type]['label'] diff --git a/examples/hetero_node_classification.py b/examples/hetero_node_classification.py new file mode 100644 index 0000000..47b0bb7 --- /dev/null +++ b/examples/hetero_node_classification.py @@ -0,0 +1,13 @@ +import os +os.environ["AUTOGL_BACKEND"] = 'dgl' + +from autogl.datasets import build_dataset_from_name +from autogl.solver import AutoHeteroNodeClassifier + +if __name__ == '__main__': + acm = build_dataset_from_name("hetero-acm-han") + solver = AutoHeteroNodeClassifier(max_evals=10) + solver.fit(acm) + acc = solver.evaluate(metric='acc') + + print("acc: ", acc) diff --git a/test/performance/heterogeneous/dgl/han_main.py b/test/performance/heterogeneous/dgl/han_main.py index 9792af7..e6d70eb 100644 --- a/test/performance/heterogeneous/dgl/han_main.py +++ b/test/performance/heterogeneous/dgl/han_main.py @@ -64,9 +64,9 @@ def main(args): num_classes = labels.max().item() + 1 labels = labels.to(args['device']) - train_mask = g.nodes[node_type].data['train_mask'].to(args['device']).bool() - val_mask = g.nodes[node_type].data['val_mask'].to(args['device']).bool() - test_mask = g.nodes[node_type].data['test_mask'].to(args['device']).bool() + train_mask = g.nodes[node_type].data['train_mask'].to(args['device']) + val_mask = g.nodes[node_type].data['val_mask'].to(args['device']) + test_mask = g.nodes[node_type].data['test_mask'].to(args['device']) model = AutoHAN( dataset=dataset, diff --git a/test/performance/heterogeneous/dgl/helper.py b/test/performance/heterogeneous/dgl/helper.py index d727943..e798b77 100644 --- a/test/performance/heterogeneous/dgl/helper.py +++ b/test/performance/heterogeneous/dgl/helper.py @@ -35,3 +35,31 @@ class EarlyStopping(object): def load_checkpoint(self, model): """Load the latest checkpoint.""" model.load_state_dict(pickle.loads(self.model)) + +def get_encoder_decoder_hp(model='han'): + if model == "han": + return { + "num_layers": 2, + "hidden": [256], ## + "heads": [8], ## + "dropout": 0.2, + "act": "gelu", + }, None + if model == "hgt": + return { + "num_layers": 2, + "hidden": [256,256,256], + "heads": 4, + "dropout": 0.2, + "act": "gelu", + "use_norm": True, + }, None + if model == "HeteroRGCN": + return { + "num_layers": 2, + "hidden": [256], + "heads": 4, + "dropout": 0.2, + "act": "leaky_relu", + }, None + return {}, None \ No newline at end of file diff --git a/test/performance/heterogeneous/dgl/solver.py b/test/performance/heterogeneous/dgl/solver.py new file mode 100644 index 0000000..02e6868 --- /dev/null +++ b/test/performance/heterogeneous/dgl/solver.py @@ -0,0 +1,59 @@ +import numpy as np +from autogl.datasets import build_dataset_from_name +from autogl.solver import AutoHeteroNodeClassifier +from helper import get_encoder_decoder_hp +from tqdm import tqdm + +def fixed(**kwargs): + return [{ + 'parameterName': k, + "type": "FIXED", + "value": v + } for k, v in kwargs.items()] + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, choices=["han", "hgt", "HeteroRGCN"], default="hgt") + parser.add_argument("--epoch", type=int, default=200) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--weight_decay", type=float, default=1e-2) + parser.add_argument("--device", type=str, default="cuda") + parser.add_argument("--repeat", type=int, default=10) + + args = parser.parse_args() + + dataset = { + "han": "hetero-acm-han", + "hgt": "hetero-acm-hgt", + "HeteroRGCN": "hetero-acm-hgt" + } + + dataset = build_dataset_from_name(dataset[args.model]) + + model_hp, _ = get_encoder_decoder_hp(args.model) + + accs = [] + process = tqdm(total=args.repeat) + for rep in range(args.repeat): + solver = AutoHeteroNodeClassifier( + graph_models=[args.model], + hpo_module="random", + ensemble_module=None, + max_evals=1, + device=args.device, + trainer_hp_space=fixed( + max_epoch=args.epoch, + early_stopping_round=args.epoch + 1, + lr=args.lr, + weight_decay=args.weight_decay + ), + model_hp_spaces=[fixed(**model_hp)] + ) + solver.fit(dataset) + acc = solver.evaluate() + accs.append(acc) + process.update(1) + process.set_postfix(mean=np.mean(accs), std=np.std(accs)) + process.close() + print("mean: {:.4f} ~ std: {:.4f}".format(np.mean(accs), np.std(accs))) diff --git a/test/performance/heterogeneous/dgl/train_hgt.py b/test/performance/heterogeneous/dgl/train_hgt.py index d35ed78..6a4d37a 100644 --- a/test/performance/heterogeneous/dgl/train_hgt.py +++ b/test/performance/heterogeneous/dgl/train_hgt.py @@ -30,7 +30,7 @@ def main(args): num_features = 256 labels = G.nodes[field].data['label'].to(args.device) num_classes = labels.max().item()+1 - test_mask = G.nodes[field].data['test_mask'].bool().to(args.device) + test_mask = G.nodes[field].data['test_mask'].to(args.device) accs = [] for seed in tqdm(range(args.repeat)):