From bedd87aca4c2a83bb1721d797fa83db1230965ef Mon Sep 17 00:00:00 2001 From: lihy96 Date: Thu, 9 Sep 2021 16:46:46 +0800 Subject: [PATCH] merge dgl into trainer --- autogl/module/model/__init__.py | 24 ++--- autogl/module/train/__init__.py | 66 ++++++------- autogl/module/train/base.py | 4 +- .../module/train/graph_classification_full.py | 92 ++++++++++++++----- 4 files changed, 116 insertions(+), 70 deletions(-) diff --git a/autogl/module/model/__init__.py b/autogl/module/model/__init__.py index c995d9c..9763f83 100644 --- a/autogl/module/model/__init__.py +++ b/autogl/module/model/__init__.py @@ -1,12 +1,12 @@ -import importlib -import sys -from ...backend import DependentBackend - -# load corresponding backend of subclass -def _load_subclass_backend(backend): - sub_module = importlib.import_module(f'.{backend.get_backend_name()}', __name__) - this = sys.modules[__name__] - for api, obj in sub_module.__dict__.items(): - setattr(this, api, obj) - -_load_subclass_backend(DependentBackend) +# import importlib +# import sys +# from ...backend import DependentBackend +# +# # load corresponding backend of subclass +# def _load_subclass_backend(backend): +# sub_module = importlib.import_module(f'.{backend.get_backend_name()}', __name__) +# this = sys.modules[__name__] +# for api, obj in sub_module.__dict__.items(): +# setattr(this, api, obj) +# +# _load_subclass_backend(DependentBackend) diff --git a/autogl/module/train/__init__.py b/autogl/module/train/__init__.py index cacdc8d..400d9bc 100644 --- a/autogl/module/train/__init__.py +++ b/autogl/module/train/__init__.py @@ -1,13 +1,13 @@ TRAINER_DICT = {} from .base import ( BaseTrainer, - Evaluation, - BaseNodeClassificationTrainer, - BaseGraphClassificationTrainer, - BaseLinkPredictionTrainer, +# Evaluation, +# BaseNodeClassificationTrainer, +# BaseGraphClassificationTrainer, +# BaseLinkPredictionTrainer, ) - - +# +# def register_trainer(name): def register_trainer_cls(cls): if name in TRAINER_DICT: @@ -20,30 +20,30 @@ def register_trainer(name): return cls return register_trainer_cls - - -from .graph_classification_full import GraphClassificationFullTrainer -from .node_classification_full import NodeClassificationFullTrainer -from .link_prediction import LinkPredictionTrainer -from .node_classification_trainer import * -from .evaluation import get_feval, Acc, Auc, Logloss, Mrr, MicroF1 - -__all__ = [ - "BaseTrainer", - "Evaluation", - "BaseGraphClassificationTrainer", - "BaseNodeClassificationTrainer", - "BaseLinkPredictionTrainer", - "GraphClassificationFullTrainer", - "NodeClassificationFullTrainer", - "NodeClassificationGraphSAINTTrainer", - "NodeClassificationLayerDependentImportanceSamplingTrainer", - "NodeClassificationNeighborSamplingTrainer", - "LinkPredictionTrainer", - "Acc", - "Auc", - "Logloss", - "Mrr", - "MicroF1", - "get_feval", -] +# +# +# from .graph_classification_full import GraphClassificationFullTrainer +# from .node_classification_full import NodeClassificationFullTrainer +# from .link_prediction import LinkPredictionTrainer +# from .node_classification_trainer import * +# from .evaluation import get_feval, Acc, Auc, Logloss, Mrr, MicroF1 +# +# __all__ = [ +# "BaseTrainer", +# "Evaluation", +# "BaseGraphClassificationTrainer", +# "BaseNodeClassificationTrainer", +# "BaseLinkPredictionTrainer", +# "GraphClassificationFullTrainer", +# "NodeClassificationFullTrainer", +# "NodeClassificationGraphSAINTTrainer", +# "NodeClassificationLayerDependentImportanceSamplingTrainer", +# "NodeClassificationNeighborSamplingTrainer", +# "LinkPredictionTrainer", +# "Acc", +# "Auc", +# "Logloss", +# "Mrr", +# "MicroF1", +# "get_feval", +# ] diff --git a/autogl/module/train/base.py b/autogl/module/train/base.py index 1797bb5..1e6b1fb 100644 --- a/autogl/module/train/base.py +++ b/autogl/module/train/base.py @@ -3,7 +3,9 @@ import typing as _typing import torch import pickle -from ..model import BaseModel, ModelUniversalRegistry +# from ..model import BaseModel, ModelUniversalRegistry +from ..model.dgl.base import BaseModel +from ..model.dgl._model_registry import ModelUniversalRegistry from .evaluation import Evaluation, get_feval, Acc from ...utils import get_logger diff --git a/autogl/module/train/graph_classification_full.py b/autogl/module/train/graph_classification_full.py index 8b12a05..6ecb015 100644 --- a/autogl/module/train/graph_classification_full.py +++ b/autogl/module/train/graph_classification_full.py @@ -8,7 +8,8 @@ from torch.optim.lr_scheduler import ( ReduceLROnPlateau, ) import torch.nn.functional as F -from ..model import BaseModel +# from ..model import BaseModel +from ..model.dgl.base import BaseModel from .evaluation import get_feval, Logloss from typing import Union from ...datasets import utils @@ -71,6 +72,7 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): feval=[Logloss], loss="nll_loss", lr_scheduler_type=None, + criterion=None, *args, **kwargs ): @@ -124,6 +126,9 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): self.initialized = False self.device = device + self.pyg_dgl = 'dgl' + self.criterion = criterion + self.space = [ { "parameterName": "max_epoch", @@ -233,18 +238,31 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): self.model.model.train() loss_all = 0 for data in train_loader: - data = data.to(self.device) - optimizer.zero_grad() - output = self.model.model(data) - # loss = F.nll_loss(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) - ) - loss.backward() - loss_all += data.num_graphs * loss.item() + if self.pyg_dgl == 'pyg': + data = data.to(self.device) + optimizer.zero_grad() + output = self.model.model(data) + # loss = F.nll_loss(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) + ) + loss.backward() + loss_all += data.num_graphs * loss.item() + elif self.pyg_dgl == 'dgl': + graphs, labels = data + graphs = graphs.to(self.device) + labels = labels.to(self.device) + feat = graphs.ndata.pop('attr') + output = self.model.model(graphs, feat) + loss = self.criterion(output, labels) + + loss.backward() + + loss_all += len(labels) * loss.item() + optimizer.step() if self.lr_scheduler_type: scheduler.step() @@ -259,9 +277,13 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): if eval_func.is_higher_better(): val_loss = -val_loss self.early_stopping(val_loss, self.model.model) + + print('val_loss', val_loss) + 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) @@ -282,8 +304,17 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): self.model.model.eval() pred = [] for data in loader: - data = data.to(self.device) - pred.append(self.model.model(data)) + if self.pyg_dgl == 'pyg': + data = data.to(self.device) + pred.append(self.model.model(data)) + elif self.pyg_dgl == 'dgl': + graphs, labels = data + graphs = graphs.to(self.device) + labels = labels.to(self.device) + feat = graphs.ndata.pop('attr') + outputs = self.model.model(graphs, feat) + pred.append(outputs) + ret = torch.cat(pred, 0) return ret @@ -332,10 +363,15 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): ------- The prediction result of ``predict_proba``. """ - loader = utils.graph_get_split( - dataset, mask, batch_size=self.batch_size, num_workers=self.num_workers - ) - return self._predict_proba(loader, in_log_format=True).max(1)[1] + if self.pyg_dgl == 'pyg': + loader = utils.graph_get_split( + dataset, mask, batch_size=self.batch_size, num_workers=self.num_workers + ) + return self._predict_proba(loader, in_log_format=True).max(1)[1] + + elif self.pyg_dgl == 'dgl': + loader = dataset + return self._predict_proba(loader, in_log_format=True).max(1)[1] def predict_proba(self, dataset, mask="test", in_log_format=False): """ @@ -430,10 +466,14 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): res: The evaluation result on the given dataset. """ - loader = utils.graph_get_split( - dataset, mask, batch_size=self.batch_size, num_workers=self.num_workers - ) - return self._evaluate(loader, feval) + if self.pyg_dgl == 'pyg': + loader = utils.graph_get_split( + dataset, mask, batch_size=self.batch_size, num_workers=self.num_workers + ) + return self._evaluate(loader, feval) + elif self.pyg_dgl == 'dgl': + loader = dataset + return self._evaluate(loader, feval) def _evaluate(self, loader, feval=None): if feval is None: @@ -445,7 +485,11 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): y_true_tmp = [] for data in loader: - y_true_tmp.append(data.y) + if self.pyg_dgl == 'pyg': + y_true_tmp.append(data.y) + elif self.pyg_dgl == 'dgl': + graphs, labels = data + y_true_tmp.append(labels) y_true = torch.cat(y_true_tmp, 0) if not isinstance(feval, list):