GraphGym Test Suite (pyg-team#3216)

* add basic graphgym test

* tensorboard

* typo

* update

* add print

* fix windows issues

test/graphgym/example_node.yml

@@ -0,0 +1,37 @@
+tensorboard_each_run: False
+tensorboard_agg: False
+dataset:
+  format: PyG
+  name: Cora
+  task: node
+  task_type: classification
+  node_encoder: False
+  node_encoder_name: Atom
+  edge_encoder: False
+  edge_encoder_name: Bond
+train:
+  batch_size: 128
+  eval_period: 1
+  ckpt_period: 100
+  sampler: full_batch
+model:
+  type: gnn
+  loss_fun: cross_entropy
+  edge_decoding: dot
+  graph_pooling: add
+gnn:
+  layers_pre_mp: 0
+  layers_mp: 2
+  layers_post_mp: 1
+  dim_inner: 16
+  layer_type: gcnconv
+  stage_type: stack
+  batchnorm: False
+  act: prelu
+  dropout: 0.1
+  agg: mean
+  normalize_adj: False
+optim:
+  optimizer: adam
+  base_lr: 0.01
+  max_epoch: 3

test/graphgym/test_graphgym.py

@@ -0,0 +1,76 @@
+import os.path as osp
+import torch
+import random
+import sys
+import shutil
+
+from collections import namedtuple
+
+from torch_geometric import seed_everything
+from torch_geometric.graphgym.train import train
+from torch_geometric.graphgym.loader import create_loader
+from torch_geometric.graphgym.model_builder import create_model
+from torch_geometric.graphgym.models.head import GNNNodeHead
+from torch_geometric.graphgym.logger import set_printing, create_logger
+from torch_geometric.graphgym.models.gnn import FeatureEncoder, GNNStackStage
+from torch_geometric.graphgym.config import (cfg, dump_cfg, set_run_dir,
+                                             set_agg_dir, load_cfg)
+from torch_geometric.graphgym.utils import (agg_runs, params_count,
+                                            auto_select_device)
+from torch_geometric.graphgym.optimizer import (create_optimizer,
+                                                create_scheduler,
+                                                OptimizerConfig,
+                                                SchedulerConfig)
+
+
+def test_run_single_graphgym():
+    Args = namedtuple('Args', ['cfg_file', 'opts'])
+    root = osp.join(osp.dirname(osp.realpath(__file__)))
+    args = Args(osp.join(root, 'example_node.yml'), [])
+
+    load_cfg(cfg, args)
+    cfg.out_dir = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
+    cfg.run_dir = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
+    cfg.dataset.dir = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
+    dump_cfg(cfg)
+    set_printing()
+
+    seed_everything(cfg.seed)
+    auto_select_device()
+    set_run_dir(cfg.out_dir, args.cfg_file)
+
+    loaders = create_loader()
+    assert len(loaders) == 3
+
+    loggers = create_logger()
+    assert len(loggers) == 3
+
+    model = create_model()
+    assert isinstance(model, torch.nn.Module)
+    assert isinstance(model.encoder, FeatureEncoder)
+    assert isinstance(model.mp, GNNStackStage)
+    assert isinstance(model.post_mp, GNNNodeHead)
+
+    optimizer_config = OptimizerConfig(optimizer=cfg.optim.optimizer,
+                                       base_lr=cfg.optim.base_lr,
+                                       weight_decay=cfg.optim.weight_decay,
+                                       momentum=cfg.optim.momentum)
+    optimizer = create_optimizer(model.parameters(), optimizer_config)
+    assert isinstance(optimizer, torch.optim.Adam)
+
+    scheduler_config = SchedulerConfig(scheduler=cfg.optim.scheduler,
+                                       steps=cfg.optim.steps,
+                                       lr_decay=cfg.optim.lr_decay,
+                                       max_epoch=cfg.optim.max_epoch)
+    scheduler = create_scheduler(optimizer, scheduler_config)
+    assert isinstance(scheduler, torch.optim.lr_scheduler.CosineAnnealingLR)
+
+    cfg.params = params_count(model)
+    assert cfg.params == 23336
+
+    train(loggers, loaders, model, optimizer, scheduler)
+
+    agg_runs(set_agg_dir(cfg.out_dir, args.cfg_file), cfg.metric_best)
+
+    shutil.rmtree(cfg.out_dir)
+    shutil.rmtree(cfg.dataset.dir)

torch_geometric/graphgym/config.py

@@ -3,6 +3,7 @@
 from yacs.config import CfgNode as CN
 import shutil
 
+from torch_geometric.data.makedirs import makedirs
 import torch_geometric.graphgym.register as register
 
 # Global config object
@@ -482,6 +483,7 @@ def dump_cfg(cfg):
         cfg (CfgNode): Configuration node
 
     """
+    makedirs(cfg.out_dir)
     cfg_file = os.path.join(cfg.out_dir, cfg.cfg_dest)
     with open(cfg_file, 'w') as f:
         cfg.dump(stream=f)
@@ -519,6 +521,8 @@ def set_run_dir(out_dir, fname):
     fname = fname.split('/')[-1]
     if fname.endswith('.yaml'):
         fname = fname[:-5]
+    elif fname.endswith('.yml'):
+        fname = fname[:-4]
     cfg.run_dir = os.path.join(out_dir, fname, str(cfg.seed))
     # Make output directory
     if cfg.train.auto_resume:
@@ -540,6 +544,8 @@ def set_agg_dir(out_dir, fname):
     fname = fname.split('/')[-1]
     if fname.endswith('.yaml'):
         fname = fname[:-5]
+    elif fname.endswith('.yml'):
+        fname = fname[:-4]
     return os.path.join(out_dir, fname)
 
 

torch_geometric/graphgym/logger.py

@@ -1,21 +1,17 @@
 import torch
 import math
-import os
 import sys
 import logging
+
 from torch_geometric.graphgym.config import cfg
+from torch_geometric.data.makedirs import makedirs
 from torch_geometric.graphgym.utils.io import dict_to_json, dict_to_tb
 
 from sklearn.metrics import accuracy_score, precision_score, recall_score, \
     f1_score, roc_auc_score, mean_absolute_error, mean_squared_error
 
 from torch_geometric.graphgym.utils.device import get_current_gpu_usage
 
-try:
-    from tensorboardX import SummaryWriter
-except ImportError:
-    SummaryWriter = None
-
 
 def set_printing():
     """
@@ -24,6 +20,7 @@ def set_printing():
     """
     logging.root.handlers = []
     logging_cfg = {'level': logging.INFO, 'format': '%(message)s'}
+    makedirs(cfg.run_dir)
     h_file = logging.FileHandler('{}/logging.log'.format(cfg.run_dir))
     h_stdout = logging.StreamHandler(sys.stdout)
     if cfg.print == 'file':
@@ -46,11 +43,9 @@ def __init__(self, name='train', task_type=None):
         self._time_total = 0  # won't be reset
 
         self.out_dir = '{}/{}'.format(cfg.run_dir, name)
-        os.makedirs(self.out_dir, exist_ok=True)
+        makedirs(self.out_dir)
         if cfg.tensorboard_each_run:
-            if SummaryWriter is None:
-                raise ImportError(
-                    'Tensorboard support requires `tensorboardX`.')
+            from tensorboardX import SummaryWriter
             self.tb_writer = SummaryWriter(self.out_dir)
 
         self.reset()