chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,10 @@
|
||||
try:
|
||||
from .base import BaseTask
|
||||
from .hypergraphs import HypergraphVertexClassificationTask
|
||||
|
||||
|
||||
except:
|
||||
print(
|
||||
"Warning raise in module: experiments. Please install Pytorch before you use"
|
||||
" functions related to nueral network"
|
||||
)
|
||||
@@ -0,0 +1,204 @@
|
||||
import abc
|
||||
import logging
|
||||
import shutil
|
||||
import time
|
||||
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
from typing import Optional
|
||||
from typing import Union
|
||||
|
||||
import optuna
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from easygraph.classes.base import load_structure
|
||||
from easygraph.ml_metrics import BaseEvaluator
|
||||
from easygraph.utils import default_log_formatter
|
||||
from optuna.samplers import TPESampler
|
||||
|
||||
|
||||
class BaseTask:
|
||||
r"""The base class of Auto-experiment in EasyGraph.
|
||||
|
||||
Args:
|
||||
``work_root`` (``Optional[Union[str, Path]]``): User's work root to store all studies.
|
||||
``data`` (``dict``): The dictionary to store input data that used in the experiment.
|
||||
``model_builder`` (``Callable``): The function to build a model with a fixed parameter ``trial``.
|
||||
``train_builder`` (``Callable``): The function to build a training configuration with two fixed parameters ``trial`` and ``model``.
|
||||
``evaluator`` (``eg.ml_metrics.BaseEvaluator``): The EasyGraph evaluator object to evaluate performance of the model in the experiment.
|
||||
``device`` (``torch.device``): The target device to run the experiment.
|
||||
``structure_builder`` (``Optional[Callable]``): The function to build a structure with a fixed parameter ``trial``. The structure can be ``eg.Graph``, ``eg.DiGraph``, ``eg.BiGraph``, and ``eg.Hypergraph``.
|
||||
``study_name`` (``Optional[str]``): The name of this study. If set to ``None``, the study name will be generated automatically according to current time. Defaults to ``None``.
|
||||
``overwrite`` (``bool``): The flag that whether to overwrite the existing study. Different studies are identified by the ``study_name``. Defaults to ``True``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
work_root: Optional[Union[str, Path]],
|
||||
data: dict,
|
||||
model_builder: Callable,
|
||||
train_builder: Callable,
|
||||
evaluator: BaseEvaluator,
|
||||
device: torch.device,
|
||||
structure_builder: Optional[Callable] = None,
|
||||
study_name: Optional[str] = None,
|
||||
overwrite: bool = True,
|
||||
):
|
||||
self.data = data
|
||||
self.model_builder = model_builder
|
||||
self.train_builder = train_builder
|
||||
self.structure_builder = structure_builder
|
||||
self.evaluator = evaluator
|
||||
self.device = device
|
||||
self.study = None
|
||||
if study_name is None:
|
||||
self.study_name = time.strftime("%Y-%m-%d--%H-%M-%S", time.localtime())
|
||||
else:
|
||||
self.study_name = study_name
|
||||
work_root = Path(work_root)
|
||||
self.study_root = work_root / self.study_name
|
||||
if overwrite and self.study_root.exists():
|
||||
shutil.rmtree(self.study_root)
|
||||
self.log_file = self.study_root / "log.txt"
|
||||
self.cache_root = self.study_root / "cache"
|
||||
if not work_root.exists():
|
||||
if work_root.parent.exists():
|
||||
work_root.mkdir(exist_ok=True)
|
||||
else:
|
||||
raise ValueError(f"The work_root {work_root} does not exist.")
|
||||
self.study_root.mkdir(exist_ok=True)
|
||||
self.cache_root.mkdir(exist_ok=True)
|
||||
# configure logging
|
||||
self.logger = optuna.logging.get_logger("optuna")
|
||||
self.logger.setLevel(logging.INFO)
|
||||
out_file_handler = logging.FileHandler(self.log_file, mode="a", encoding="utf8")
|
||||
out_file_handler.setFormatter(default_log_formatter())
|
||||
self.logger.addHandler(out_file_handler)
|
||||
self.logger.info(f"Logs will be saved to {self.log_file.absolute()}")
|
||||
self.logger.info(
|
||||
f"Files in training will be saved in {self.study_root.absolute()}"
|
||||
)
|
||||
|
||||
def experiment(self, trial: optuna.Trial):
|
||||
r"""Run the experiment for a given trial.
|
||||
|
||||
Args:
|
||||
``trial`` (``optuna.Trial``): The ``optuna.Trial`` object.
|
||||
"""
|
||||
if self.structure_builder is not None:
|
||||
self.data["structure"] = self.structure_builder(trial).to(self.device)
|
||||
model = self.model_builder(trial).to(self.device)
|
||||
train_configs: dict = self.train_builder(trial, model)
|
||||
assert "optimizer" in train_configs.keys()
|
||||
optimizer = train_configs["optimizer"]
|
||||
assert "criterion" in train_configs.keys()
|
||||
criterion = train_configs["criterion"]
|
||||
scheduler = train_configs.get("scheduler", None)
|
||||
|
||||
best_model = None
|
||||
if self.direction == "maximize":
|
||||
best_score = -float("inf")
|
||||
else:
|
||||
best_score = float("inf")
|
||||
for epoch in range(self.max_epoch):
|
||||
self.train(self.data, model, optimizer, criterion)
|
||||
val_res = self.validate(self.data, model)
|
||||
trial.report(val_res, epoch)
|
||||
if trial.should_prune():
|
||||
raise optuna.exceptions.TrialPruned()
|
||||
if scheduler is not None:
|
||||
scheduler.step()
|
||||
if self.direction == "maximize":
|
||||
if val_res > best_score:
|
||||
best_score = val_res
|
||||
best_model = deepcopy(model)
|
||||
with open(self.cache_root / f"{trial.number}_model.pth", "wb") as f:
|
||||
torch.save(best_model.cpu().state_dict(), f)
|
||||
self.data["structure"].save(self.cache_root / f"{trial.number}_structure.dhg")
|
||||
return best_score
|
||||
|
||||
def _remove_cached_data(self):
|
||||
r"""Remove cached models and structures."""
|
||||
if self.study is not None:
|
||||
for filename in self.cache_root.glob("*"):
|
||||
if filename.stem.split("_")[0] != str(self.study.best_trial.number):
|
||||
filename.unlink()
|
||||
|
||||
def run(self, max_epoch: int, num_trials: int = 1, direction: str = "maximize"):
|
||||
r"""Run experiments with automatically hyper-parameter tuning.
|
||||
|
||||
Args:
|
||||
``max_epoch`` (``int``): The maximum number of epochs to train for each experiment.
|
||||
``num_trials`` (``int``): The number of trials to run. Defaults to ``1``.
|
||||
``direction`` (``str``): The direction to optimize. Defaults to ``"maximize"``.
|
||||
"""
|
||||
self.logger.info(f"Random seed is {dhg.random.seed()}")
|
||||
sampler = TPESampler(seed=dhg.random.seed())
|
||||
self.max_epoch, self.direction = max_epoch, direction
|
||||
self.study = optuna.create_study(direction=direction, sampler=sampler)
|
||||
self.study.optimize(self.experiment, n_trials=num_trials, timeout=600)
|
||||
|
||||
self._remove_cached_data()
|
||||
self.best_model = self.model_builder(self.study.best_trial)
|
||||
self.best_model.load_state_dict(
|
||||
torch.load(f"{self.cache_root}/{self.study.best_trial.number}_model.pth")
|
||||
)
|
||||
self.best_structure = load_structure(
|
||||
f"{self.cache_root}/{self.study.best_trial.number}_structure.dhg"
|
||||
)
|
||||
self.best_model = self.best_model.to(self.device)
|
||||
self.best_structure = self.best_structure.to(self.device)
|
||||
|
||||
self.logger.info("Best trial:")
|
||||
self.best_trial = self.study.best_trial
|
||||
self.logger.info(f"\tValue: {self.best_trial.value:.3f}")
|
||||
self.logger.info(f"\tParams:")
|
||||
for key, value in self.best_trial.params.items():
|
||||
self.logger.info(f"\t\t{key} |-> {value}")
|
||||
test_res = self.test()
|
||||
self.logger.info(f"Final test results:")
|
||||
for key, value in test_res.items():
|
||||
self.logger.info(f"\t{key} |-> {value:.3f}")
|
||||
|
||||
@abc.abstractmethod
|
||||
def train(
|
||||
self,
|
||||
data: dict,
|
||||
model: nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
criterion: nn.Module,
|
||||
):
|
||||
r"""Train model for one epoch.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``): The input data.
|
||||
``model`` (``nn.Module``): The model.
|
||||
``optimizer`` (``torch.optim.Optimizer``): The model optimizer.
|
||||
``criterion`` (``nn.Module``): The loss function.
|
||||
"""
|
||||
|
||||
@torch.no_grad()
|
||||
@abc.abstractmethod
|
||||
def validate(
|
||||
self,
|
||||
data: dict,
|
||||
model: nn.Module,
|
||||
):
|
||||
r"""Validate the model.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``): The input data.
|
||||
``model`` (``nn.Module``): The model.
|
||||
"""
|
||||
|
||||
@torch.no_grad()
|
||||
@abc.abstractmethod
|
||||
def test(self, data: Optional[dict] = None, model: Optional[nn.Module] = None):
|
||||
r"""Test the model.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``, optional): The input data if set to ``None``, the specified ``data`` in the initialization of the experiments will be used. Defaults to ``None``.
|
||||
``model`` (``nn.Module``, optional): The model if set to ``None``, the trained best model will be used. Defaults to ``None``.
|
||||
"""
|
||||
@@ -0,0 +1 @@
|
||||
from .hypergraph import HypergraphVertexClassificationTask
|
||||
@@ -0,0 +1,121 @@
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
from typing import Optional
|
||||
from typing import Union
|
||||
|
||||
import optuna
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from easygraph.ml_metrics import BaseEvaluator
|
||||
|
||||
from ..vertex_classification import VertexClassificationTask
|
||||
|
||||
|
||||
class HypergraphVertexClassificationTask(VertexClassificationTask):
|
||||
r"""The auto-experiment class for the vertex classification task on hypergraph.
|
||||
|
||||
Args:
|
||||
``work_root`` (``Optional[Union[str, Path]]``): User's work root to store all studies.
|
||||
``data`` (``dict``): The dictionary to store input data that used in the experiment.
|
||||
``model_builder`` (``Callable``): The function to build a model with a fixed parameter ``trial``.
|
||||
``train_builder`` (``Callable``): The function to build a training configuration with two fixed parameters ``trial`` and ``model``.
|
||||
``evaluator`` (``easygraph.ml_metrics.BaseEvaluator``): The DHG evaluator object to evaluate performance of the model in the experiment.
|
||||
``device`` (``torch.device``): The target device to run the experiment.
|
||||
``structure_builder`` (``Optional[Callable]``): The function to build a structure with a fixed parameter ``trial``. The structure should be ``easygraph.Hypergraph``.
|
||||
``study_name`` (``Optional[str]``): The name of this study. If set to ``None``, the study name will be generated automatically according to current time. Defaults to ``None``.
|
||||
``overwrite`` (``bool``): The flag that whether to overwrite the existing study. Different studies are identified by the ``study_name``. Defaults to ``True``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
work_root: Optional[Union[str, Path]],
|
||||
data: dict,
|
||||
model_builder: Callable,
|
||||
train_builder: Callable,
|
||||
evaluator: BaseEvaluator,
|
||||
device: torch.device,
|
||||
structure_builder: Optional[Callable] = None,
|
||||
study_name: Optional[str] = None,
|
||||
overwrite: bool = True,
|
||||
):
|
||||
super().__init__(
|
||||
work_root,
|
||||
data,
|
||||
model_builder,
|
||||
train_builder,
|
||||
evaluator,
|
||||
device,
|
||||
structure_builder=structure_builder,
|
||||
study_name=study_name,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
def to(self, device: torch.device):
|
||||
r"""Move the input data to the target device.
|
||||
|
||||
Args:
|
||||
``device`` (``torch.device``): The specified target device to store the input data.
|
||||
"""
|
||||
return super().to(device)
|
||||
|
||||
@property
|
||||
def vars_for_DL(self):
|
||||
r"""Return a name list for available variables for deep learning in the vertex classification on hypergraph. The name list includes ``features``, ``structure``, ``labels``, ``train_mask``, ``val_mask``, and ``test_mask``.
|
||||
"""
|
||||
return super().vars_for_DL
|
||||
|
||||
def experiment(self, trial: optuna.Trial):
|
||||
r"""Run the experiment for a given trial.
|
||||
|
||||
Args:
|
||||
``trial`` (``optuna.Trial``): The ``optuna.Trial`` object.
|
||||
"""
|
||||
return super().experiment(trial)
|
||||
|
||||
def run(self, max_epoch: int, num_trials: int = 1, direction: str = "maximize"):
|
||||
r"""Run experiments with automatically hyper-parameter tuning.
|
||||
|
||||
Args:
|
||||
``max_epoch`` (``int``): The maximum number of epochs to train for each experiment.
|
||||
``num_trials`` (``int``): The number of trials to run. Defaults to ``1``.
|
||||
``direction`` (``str``): The direction to optimize. Defaults to ``"maximize"``.
|
||||
"""
|
||||
return super().run(max_epoch, num_trials, direction)
|
||||
|
||||
def train(
|
||||
self,
|
||||
data: dict,
|
||||
model: nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
criterion: nn.Module,
|
||||
):
|
||||
r"""Train model for one epoch.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``): The input data.
|
||||
``model`` (``nn.Module``): The model.
|
||||
``optimizer`` (``torch.optim.Optimizer``): The model optimizer.
|
||||
``criterion`` (``nn.Module``): The loss function.
|
||||
"""
|
||||
return super().train(data, model, optimizer, criterion)
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(self, data: dict, model: nn.Module):
|
||||
r"""Validate the model.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``): The input data.
|
||||
``model`` (``nn.Module``): The model.
|
||||
"""
|
||||
return super().validate(data, model)
|
||||
|
||||
@torch.no_grad()
|
||||
def test(self, data: Optional[dict] = None, model: Optional[nn.Module] = None):
|
||||
r"""Test the model.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``, optional): The input data if set to ``None``, the specified ``data`` in the intialization of the experiments will be used. Defaults to ``None``.
|
||||
``model`` (``nn.Module``, optional): The model if set to ``None``, the trained best model will be used. Defaults to ``None``.
|
||||
"""
|
||||
return super().test(data, model)
|
||||
@@ -0,0 +1,166 @@
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
from typing import Optional
|
||||
from typing import Union
|
||||
|
||||
import optuna
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from easygraph.ml_metrics import BaseEvaluator
|
||||
|
||||
from .base import BaseTask
|
||||
|
||||
|
||||
class VertexClassificationTask(BaseTask):
|
||||
r"""The auto-experiment class for the vertex classification task.
|
||||
|
||||
Args:
|
||||
``work_root`` (``Optional[Union[str, Path]]``): User's work root to store all studies.
|
||||
``data`` (``dict``): The dictionary to store input data that used in the experiment.
|
||||
``model_builder`` (``Callable``): The function to build a model with a fixed parameter ``trial``.
|
||||
``train_builder`` (``Callable``): The function to build a training configuration with two fixed parameters ``trial`` and ``model``.
|
||||
``evaluator`` (``eg.ml_metrics.BaseEvaluator``): The DHG evaluator object to evaluate performance of the model in the experiment.
|
||||
``device`` (``torch.device``): The target device to run the experiment.
|
||||
``structure_builder`` (``Optional[Callable]``): The function to build a structure with a fixed parameter ``trial``. The structure can be ``eg.Hypergraph``.
|
||||
``study_name`` (``Optional[str]``): The name of this study. If set to ``None``, the study name will be generated automatically according to current time. Defaults to ``None``.
|
||||
``overwrite`` (``bool``): The flag that whether to overwrite the existing study. Different studies are identified by the ``study_name``. Defaults to ``True``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
work_root: Optional[Union[str, Path]],
|
||||
data: dict,
|
||||
model_builder: Callable,
|
||||
train_builder: Callable,
|
||||
evaluator: BaseEvaluator,
|
||||
device: torch.device,
|
||||
structure_builder: Optional[Callable] = None,
|
||||
study_name: Optional[str] = None,
|
||||
overwrite: bool = True,
|
||||
):
|
||||
super().__init__(
|
||||
work_root,
|
||||
data,
|
||||
model_builder,
|
||||
train_builder,
|
||||
evaluator,
|
||||
device,
|
||||
structure_builder=structure_builder,
|
||||
study_name=study_name,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
self.to(self.device)
|
||||
|
||||
def to(self, device: torch.device):
|
||||
r"""Move the input data to the target device.
|
||||
|
||||
Args:
|
||||
``device`` (``torch.device``): The specified target device to store the input data.
|
||||
"""
|
||||
self.device = device
|
||||
for name in self.vars_for_DL:
|
||||
if name in self.data.keys():
|
||||
self.data[name] = self.data[name].to(device)
|
||||
return self
|
||||
|
||||
@property
|
||||
def vars_for_DL(self):
|
||||
r"""Return a name list for available variables for deep learning in the vertex classification task. The name list includes ``features``, ``structure``, ``labels``, ``train_mask``, ``val_mask``, and ``test_mask``.
|
||||
"""
|
||||
return (
|
||||
"features",
|
||||
"structure",
|
||||
"labels",
|
||||
"train_mask",
|
||||
"val_mask",
|
||||
"test_mask",
|
||||
)
|
||||
|
||||
def experiment(self, trial: optuna.Trial):
|
||||
r"""Run the experiment for a given trial.
|
||||
|
||||
Args:
|
||||
``trial`` (``optuna.Trial``): The ``optuna.Trial`` object.
|
||||
"""
|
||||
return super().experiment(trial)
|
||||
|
||||
def run(self, max_epoch: int, num_trials: int = 1, direction: str = "maximize"):
|
||||
r"""Run experiments with automatically hyper-parameter tuning.
|
||||
|
||||
Args:
|
||||
``max_epoch`` (``int``): The maximum number of epochs to train for each experiment.
|
||||
``num_trials`` (``int``): The number of trials to run. Defaults to ``1``.
|
||||
``direction`` (``str``): The direction to optimize. Defaults to ``"maximize"``.
|
||||
"""
|
||||
return super().run(max_epoch, num_trials, direction)
|
||||
|
||||
def train(
|
||||
self,
|
||||
data: dict,
|
||||
model: nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
criterion: nn.Module,
|
||||
):
|
||||
r"""Train model for one epoch.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``): The input data.
|
||||
``model`` (``nn.Module``): The model.
|
||||
``optimizer`` (``torch.optim.Optimizer``): The model optimizer.
|
||||
``criterion`` (``nn.Module``): The loss function.
|
||||
"""
|
||||
features, structure = data["features"], data["structure"]
|
||||
train_mask, labels = data["train_mask"], data["labels"]
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
outputs = model(features, structure)
|
||||
loss = criterion(
|
||||
outputs[train_mask],
|
||||
labels[train_mask],
|
||||
)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(self, data: dict, model: nn.Module):
|
||||
r"""Validate the model.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``): The input data.
|
||||
``model`` (``nn.Module``): The model.
|
||||
"""
|
||||
features, structure = data["features"], data["structure"]
|
||||
val_mask, labels = data["val_mask"], data["labels"]
|
||||
model.eval()
|
||||
outputs = model(features, structure)
|
||||
res = self.evaluator.validate(labels[val_mask], outputs[val_mask])
|
||||
return res
|
||||
|
||||
@torch.no_grad()
|
||||
def test(self, data: Optional[dict] = None, model: Optional[nn.Module] = None):
|
||||
r"""Test the model.
|
||||
|
||||
Args:
|
||||
``data`` (``dict``, optional): The input data if set to ``None``, the specified ``data`` in the initialization of the experiments will be used. Defaults to ``None``.
|
||||
``model`` (``nn.Module``, optional): The model if set to ``None``, the trained best model will be used. Defaults to ``None``.
|
||||
"""
|
||||
if data is None:
|
||||
features, structure = self.data["features"], self.best_structure
|
||||
test_mask, labels = self.data["test_mask"], self.data["labels"]
|
||||
else:
|
||||
features, structure = (
|
||||
data["features"].to(self.device),
|
||||
data["structure"].to(self.device),
|
||||
)
|
||||
test_mask, labels = (
|
||||
data["test_mask"].to(self.device),
|
||||
data["labels"].to(self.device),
|
||||
)
|
||||
if model is None:
|
||||
model = self.best_model
|
||||
model = model.to(self.device)
|
||||
model.eval()
|
||||
outputs = model(features, structure)
|
||||
res = self.evaluator.test(labels[test_mask], outputs[test_mask])
|
||||
return res
|
||||
Reference in New Issue
Block a user