Implement LightningModule & LoggerCallback in GraphGym

test/graphgym/test_graphgym.py

@@ -31,8 +31,16 @@
 )
 from torch_geometric.testing import withPackage
 
+try:
+    import pytorch_lightning as pl
+except ImportError:
+    pl = None
 num_trivial_metric_calls = 0
 
+Args = namedtuple('Args', ['cfg_file', 'opts'])
+root = osp.join(osp.dirname(osp.realpath(__file__)))
+args = Args(osp.join(root, 'example_node.yml'), [])
+
 
 def trivial_metric(true, pred, task_type):
     global num_trivial_metric_calls
@@ -110,3 +118,54 @@ def test_run_single_graphgym(auto_resume, skip_train_eval, use_trivial_metric):
     agg_runs(cfg.out_dir, cfg.metric_best)
 
     shutil.rmtree(cfg.out_dir)
+
+
+@withPackage('yacs')
+@withPackage('pytorch_lightning')
+def test_graphgym_module(tmpdir):
+    load_cfg(cfg, args)
+    cfg.out_dir = osp.join(tmpdir, str(random.randrange(sys.maxsize)))
+    cfg.run_dir = osp.join(tmpdir, str(random.randrange(sys.maxsize)))
+    cfg.dataset.dir = osp.join(tmpdir, 'pyg_test_datasets', 'Planetoid')
+
+    set_out_dir(cfg.out_dir, args.cfg_file)
+    dump_cfg(cfg)
+    set_printing()
+
+    seed_everything(cfg.seed)
+    auto_select_device()
+    set_run_dir(cfg.out_dir)
+
+    loaders = create_loader()
+    assert len(loaders) == 3
+
+    model = create_model()
+    assert isinstance(model, pl.LightningModule)
+
+    optimizer, scheduler = model.configure_optimizers()
+    assert isinstance(optimizer[0], torch.optim.Adam)
+    assert isinstance(scheduler[0], torch.optim.lr_scheduler.CosineAnnealingLR)
+
+    cfg.params = params_count(model)
+    assert cfg.params == 23880
+
+    keys = {"loss", "true", "pred", "pred_score", "step_end_time"}
+    # test training step
+    batch = next(iter(loaders[0]))
+    outputs = model.training_step(batch)
+    assert keys.issubset(outputs.keys())
+    assert torch.is_tensor(outputs["loss"])
+
+    # test validation step
+    batch = next(iter(loaders[1]))
+    outputs = model.validation_step(batch)
+    assert keys.issubset(outputs.keys())
+    assert torch.is_tensor(outputs["loss"])
+
+    # test test step
+    batch = next(iter(loaders[2]))
+    outputs = model.test_step(batch)
+    assert keys.issubset(outputs.keys())
+    assert torch.is_tensor(outputs["loss"])
+
+    shutil.rmtree(cfg.out_dir)

test/graphgym/test_logger.py

@@ -0,0 +1,11 @@
+from torch_geometric.graphgym.logger import Logger, LoggerCallback
+from torch_geometric.testing import withPackage
+
+
+@withPackage('yacs')
+@withPackage('pytorch_lightning')
+def test_logger_callback():
+    logger = LoggerCallback()
+    assert isinstance(logger.train_logger, Logger)
+    assert isinstance(logger.val_logger, Logger)
+    assert isinstance(logger.test_logger, Logger)

torch_geometric/graphgym/logger.py

@@ -1,6 +1,9 @@
 import logging
 import math
 import sys
+import time
+import warnings
+from typing import Any, Dict, Optional
 
 import torch
 
@@ -10,6 +13,16 @@
 from torch_geometric.graphgym.utils.device import get_current_gpu_usage
 from torch_geometric.graphgym.utils.io import dict_to_json, dict_to_tb
 
+try:
+    import pytorch_lightning as pl
+    from pytorch_lightning import Callback
+
+except ImportError:
+    pl = None
+    Callback = object
+    warnings.warn("Please install 'pytorch_lightning' for using the GraphGym "
+                  "experiment manager via 'pip install pytorch_lightning'")
+
 
 def set_printing():
     """
@@ -247,3 +260,105 @@ def create_logger():
     for i, dataset in enumerate(range(cfg.share.num_splits)):
         loggers.append(Logger(name=names[i], task_type=infer_task()))
     return loggers
+
+
+class LoggerCallback(Callback):
+    def __init__(self):
+        self._logger = create_logger()
+        self._train_epoch_start_time = None
+        self._val_epoch_start_time = None
+        self._test_epoch_start_time = None
+
+    def _create_stats(self, time_start, outputs: Dict[str, Any],
+                      trainer: "pl.Trainer") -> Dict:
+        true: torch.Tensor = outputs["true"]
+        pred: torch.Tensor = outputs["pred"]
+        pred_score: torch.Tensor = outputs["pred_score"]
+        loss: torch.Tensor = outputs["loss"]
+        lr = trainer.lr_scheduler_configs[0].scheduler.get_last_lr()[0]
+        stats = dict(true=true.detach().cpu(), pred=pred.detach().cpu(),
+                     pred_score=pred_score.detach().cpu(), loss=loss.item(),
+                     lr=lr, time_used=time.time() - time_start,
+                     params=cfg.params)
+
+        return stats
+
+    def on_train_epoch_start(self, trainer: "pl.Trainer",
+                             pl_module: "pl.LightningModule") -> None:
+        self._train_epoch_start_time = time.time()
+
+    def on_validation_epoch_start(self, trainer: "pl.Trainer",
+                                  pl_module: "pl.LightningModule") -> None:
+        self._val_epoch_start_time = time.time()
+
+    def on_test_epoch_start(self, trainer: "pl.Trainer",
+                            pl_module: "pl.LightningModule") -> None:
+        self._test_epoch_start_time = time.time()
+
+    def on_train_batch_end(
+        self,
+        trainer: "pl.Trainer",
+        pl_module: "pl.LightningModule",
+        outputs: Dict[str, Any],
+        batch: Any,
+        batch_idx: int,
+        unused: int = 0,
+    ) -> None:
+        stats = self._create_stats(self._train_epoch_start_time, outputs,
+                                   trainer)
+        self.train_logger.update_stats(**stats)
+
+    def on_validation_batch_end(
+        self,
+        trainer: "pl.Trainer",
+        pl_module: "pl.LightningModule",
+        outputs: Optional[Dict[str, Any]],
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        stats = self._create_stats(self._val_epoch_start_time, outputs,
+                                   trainer)
+        self.val_logger.update_stats(**stats)
+
+    def on_test_batch_end(
+        self,
+        trainer: "pl.Trainer",
+        pl_module: "pl.LightningModule",
+        outputs: Optional[Dict[str, Any]],
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        stats = self._create_stats(self._test_epoch_start_time, outputs,
+                                   trainer)
+        self.test_logger.update_stats(**stats)
+
+    def on_train_epoch_end(self, trainer: "pl.Trainer",
+                           pl_module: "pl.LightningModule") -> None:
+        self.train_logger.write_epoch(trainer.current_epoch)
+
+    def on_validation_epoch_end(self, trainer: "pl.Trainer",
+                                pl_module: "pl.LightningModule") -> None:
+        self.val_logger.write_epoch(trainer.current_epoch)
+
+    def on_test_epoch_end(self, trainer: "pl.Trainer",
+                          pl_module: "pl.LightningModule") -> None:
+        self.test_logger.write_epoch(trainer.current_epoch)
+
+    def on_epoch_end(self, trainer: "pl.Trainer",
+                     pl_module: "pl.LightningModule") -> None:
+        for logger in self._logger:
+            logger.close()
+
+    @property
+    def train_logger(self):
+        return self._logger[0]
+
+    @property
+    def val_logger(self):
+        return self._logger[1]
+
+    @property
+    def test_logger(self):
+        return self._logger[2]

torch_geometric/graphgym/model_builder.py

@@ -1,11 +1,19 @@
+import time
+import typing
 import warnings
+from typing import Any, Dict, Tuple
 
 import torch
 
 from torch_geometric.graphgym.config import cfg
+from torch_geometric.graphgym.loss import compute_loss
 from torch_geometric.graphgym.models.gnn import GNN
+from torch_geometric.graphgym.optim import create_optimizer, create_scheduler
 from torch_geometric.graphgym.register import network_dict, register_network
 
+if typing.TYPE_CHECKING:
+    from yacs.config import CfgNode
+
 try:
     from pytorch_lightning import LightningModule
 except ImportError:
@@ -17,14 +25,37 @@
 
 
 class GraphGymModule(LightningModule):
-    def __init__(self, dim_in, dim_out, cfg):
+    def __init__(self, dim_in, dim_out, cfg: "CfgNode"):
         super().__init__()
+        self.cfg = cfg
         self.model = network_dict[cfg.model.type](dim_in=dim_in,
                                                   dim_out=dim_out)
 
     def forward(self, *args, **kwargs):
         return self.model(*args, **kwargs)
 
+    def configure_optimizers(self) -> Tuple[Any, Any]:
+        optimizer = create_optimizer(self.model.parameters(), self.cfg.optim)
+        scheduler = create_scheduler(optimizer, self.cfg.optim)
+        return [optimizer], [scheduler]
+
+    def _shared_step(self, batch, split: str) -> Dict:
+        batch.split = split
+        pred, true = self.forward(batch)
+        loss, pred_score = compute_loss(pred, true)
+        step_end_time = time.time()
+        return dict(loss=loss, true=true, pred=pred, pred_score=pred_score,
+                    step_end_time=step_end_time)
+
+    def training_step(self, batch, *args, **kwargs):
+        return self._shared_step(batch, "train")
+
+    def validation_step(self, batch, *args, **kwargs):
+        return self._shared_step(batch, "val")
+
+    def test_step(self, batch, *args, **kwargs):
+        return self._shared_step(batch, "test")
+
     @property
     def encoder(self) -> torch.nn.Module:
         return self.model.encoder