Disentangle cfg from HeteroGNN (rusty1s#174)

* disentangle config from HeteroPostLayer

* update

* update

* update

* add doc

* fix imports

* update test

* fix API test

* fix test

* fix tests

* add doc

* MLPNodeHead

* update pyg

* reset

* resolve comments

* update doc

* typo

* typo

* fix test

* update

bin/kumo-train

@@ -16,9 +16,8 @@ from torch_geometric.graphgym.utils.device import auto_select_device
 
 import kumo.custom  # noqa
 from kumo.config import cfg
-from kumo.model import create_model
+from kumo.model.model_builder import create_model, get_emb_size
 from kumo.train import create_loader
-from kumo.train.encoder import get_emb_size
 from kumo.train.trainer import ProgressBar
 from kumo.utils import tracing
 from kumo.utils.visualization import visualize_scalar_distribution
@@ -32,49 +31,64 @@ def main(args):
     cfg_name = get_fname(args.cfg_file)
     tb_logger = pl_loggers.TensorBoardLogger(cfg.out_dir, name=cfg_name,
                                              sub_dir=str(cfg.seed + 1))
-    # Repeat for different random seeds
-    for i in range(args.repeat):
+
+    for i in range(args.repeat):  # Repeat for different random seeds:
         # Set configurations for each run
         cfg.seed = cfg.seed + 1
         seed_everything(cfg.seed)
         auto_select_device()
+
         # Logger: create sub-directory and logger for each random seed run
-        tb_logger = pl_loggers.TensorBoardLogger(cfg.out_dir, name=cfg_name,
-                                                 version=tb_logger.version,
-                                                 sub_dir=str(cfg.seed))
+        tb_logger = pl_loggers.TensorBoardLogger(
+            cfg.out_dir,
+            name=cfg_name,
+            version=tb_logger.version,
+            sub_dir=str(cfg.seed),
+        )
         ml_logger = pl_loggers.MLFlowLogger(
             experiment_name=f'{cfg_name}_v{tb_logger.version}',
-            save_dir=osp.join(cfg.out_dir,
-                              'mlruns'), artifact_location=tb_logger.log_dir)
+            save_dir=osp.join(cfg.out_dir, 'mlruns'),
+            artifact_location=tb_logger.log_dir,
+        )
         cfg.run_dir = tb_logger.log_dir
         set_printing()
+
         with open(osp.join(cfg.run_dir, cfg.cfg_dest), 'w') as f:
             cfg.dump(stream=f)
 
-        # Set machine learning pipeline
-        dataset, loaders = create_loader(cfg)
-        emb_size = get_emb_size(dataset)
-        model = create_model(cfg, dataset.metadata(), emb_size=emb_size)
-        # infer dim
-        model(dataset.feat_dict, dataset.edge_index_dict,
-              dataset.discrete_feat_dict)
+        # Set machine learning pipeline:
+        data, loaders = create_loader(cfg)
+        emb_size = get_emb_size(data)
+        model = create_model(cfg, data.metadata(), emb_size)
+
+        # Lazy initialization:
+        model(data.feat_dict, data.edge_index_dict, data.discrete_feat_dict)
 
         # Training
         gpus = 1 if cfg.device != 'cpu' and torch.cuda.is_available() else 0
+
+        callbacks = []
+
+        callbacks.append(ProgressBar())
+
+        callbacks.append(
+            ModelCheckpoint(
+                save_top_k=1,
+                monitor='val_loss',
+                mode='min',
+            ))
+
         if cfg.optim.early_stopping:
             patience = cfg.optim.patience
             if patience is None:
                 patience = cfg.optim.max_epoch // 5
-            early_stopping = EarlyStopping(monitor="val_loss",
-                                           min_delta=cfg.optim.min_delta,
-                                           patience=patience, mode="min")
-            ckpt = ModelCheckpoint(save_top_k=1, monitor='val_loss',
-                                   mode='min')
-            callbacks = [early_stopping, ckpt]
-        else:
-            callbacks = []
-
-        callbacks.append(ProgressBar())
+            callbacks.append(
+                EarlyStopping(
+                    monitor="val_loss",
+                    min_delta=cfg.optim.min_delta,
+                    patience=patience,
+                    mode="min",
+                ))
 
         trainer = pl.Trainer(
             gpus=gpus,
@@ -83,13 +97,14 @@ def main(args):
             logger=[tb_logger, ml_logger],
             callbacks=callbacks,
         )
-        # visualization
+
         visualize_scalar_distribution(
-            dataset[cfg.dataset.target_table].y,
+            data[cfg.dataset.target_table].y,
             os.path.join(tb_logger.log_dir, "target_distribution.png"),
         )
+
         trainer.fit(model, loaders[0], loaders[1])
-        # test
+
         print('Validation on best epoch:')
         trainer.validate(ckpt_path='best', dataloaders=loaders[1])
 

kumo/__init__.py

@@ -12,7 +12,7 @@
 from kumo.connector import CSVConnector, SnowflakeConnector
 from kumo.store import Store
 from kumo.task import Task
-from kumo.model import HeteroGNN
+from kumo.model import GeneralHeteroGNN
 from kumo.train import Trainer
 from kumo.run import Runner
 
@@ -33,7 +33,7 @@
     'SnowflakeConnector',
     'Store',
     'Task',
-    'HeteroGNN',
+    'GeneralHeteroGNN',
     'Trainer',
     'Runner',
     '__version__',

kumo/model/__init__.py

@@ -1,9 +1,11 @@
-from .layer import HeteroPostLayer
-from .gnn import HeteroGNN
+from .encoder import HeteroFeatureEncoder
+from .general_hetero_gnn import GeneralHeteroGNN
+from .head import MLPNodeHead
 from .model_builder import create_model
 
 __all__ = classes = [
-    'HeteroPostLayer',
-    'HeteroGNN',
-    'create_model'
+    'HeteroFeatureEncoder',
+    'GeneralHeteroGNN',
+    'MLPNodeHead',
+    'create_model',
 ]

kumo/model/encoder.py

@@ -0,0 +1,65 @@
+from typing import Dict, List
+
+import torch
+from torch import Tensor
+from torch.nn import Embedding
+from torch_geometric.nn import Linear
+from torch_geometric.typing import NodeType
+
+
+class HeteroFeatureEncoder(torch.nn.Module):
+    r"""Encodes continuous (:obj:`feat`) and discrete features
+    (:obj:`discrete_feat`) into a shared embedding space.
+
+    Args:
+        out_channels (int): Size of each output sample.
+        emb_size (Dict[str, List[int]]): The number of embeddings for each
+            discrete feature in each node type.
+        bias (bool, optional): If set to :obj:`False`, the layer will not learn
+            an additive bias. (default: :obj:`True`)
+    """
+    def __init__(
+        self,
+        out_channels: int,
+        emb_size: Dict[str, List[int]],
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.out_channels = out_channels
+
+        self.lin_dict = torch.nn.ModuleDict()
+        for node_type in emb_size.keys():
+            self.lin_dict[node_type] = Linear(-1, out_channels, bias=bias)
+
+        self.emb_dict = torch.nn.ModuleDict()
+        for node_type, sizes in emb_size.items():
+            self.emb_dict[node_type] = torch.nn.ModuleList(
+                [Embedding(size, out_channels) for size in sizes])
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for lin in self.lin_dict.values():
+            lin.reset_parameters()
+        for emb_list in self.emb_dict.values():
+            for emb in emb_list:
+                torch.nn.init.xavier_uniform_(emb.weight)
+
+    def forward(
+        self,
+        feat_dict: Dict[NodeType, Tensor],
+        discrete_feat_dict: Dict[NodeType, Tensor],
+    ) -> Dict[NodeType, Tensor]:
+        """"""
+        out_dict = {}
+        for node_type in feat_dict:
+            out = self.lin_dict[node_type](feat_dict[node_type])
+            for i, emb in enumerate(self.emb_dict[node_type]):
+                index = discrete_feat_dict[node_type][:, i]
+                out = out + emb(index)
+            out_dict[node_type] = out
+        return out_dict
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.out_channels})'