From 1b1530e4ddf66803bd0b7adf33d8d624919d02ab Mon Sep 17 00:00:00 2001 From: Marc Brockschmidt Date: Tue, 8 Oct 2019 17:46:35 +0000 Subject: [PATCH] feat(*): Support GNN-Edge-MLP0/MLP1 model names --- README.md | 17 ++++++------ models/ggnn_model.py | 2 +- models/gnn_edge_mlp_model.py | 4 +-- models/gnn_film_model.py | 2 +- models/rgat_model.py | 2 +- models/rgcn_model.py | 2 +- models/rgdcn_model.py | 2 +- models/sparse_graph_model.py | 4 +-- run_ppi_benchs.py | 2 +- tasks/default_hypers/PPI_GNN-Edge-MLP0.json | 7 +++++ ...N-Edge-MLP.json => PPI_GNN-Edge-MLP1.json} | 0 tasks/default_hypers/QM9_GNN-Edge-MLP.json | 7 ----- tasks/default_hypers/QM9_GNN-Edge-MLP0.json | 27 +++++++++++++++++++ tasks/default_hypers/QM9_GNN-Edge-MLP1.json | 27 +++++++++++++++++++ train.py | 7 ++--- utils/model_utils.py | 26 +++++++++++------- 16 files changed, 100 insertions(+), 38 deletions(-) create mode 100644 tasks/default_hypers/PPI_GNN-Edge-MLP0.json rename tasks/default_hypers/{PPI_GNN-Edge-MLP.json => PPI_GNN-Edge-MLP1.json} (100%) delete mode 100644 tasks/default_hypers/QM9_GNN-Edge-MLP.json create mode 100644 tasks/default_hypers/QM9_GNN-Edge-MLP0.json create mode 100644 tasks/default_hypers/QM9_GNN-Edge-MLP1.json diff --git a/README.md b/README.md index 97d8272..3e800e5 100755 --- a/README.md +++ b/README.md @@ -87,7 +87,7 @@ Currently, five model types are implemented: * `GGNN`: Gated Graph Neural Networks ([Li et al., 2015](#li-et-al-2015)). * `RGCN`: Relational Graph Convolutional Networks ([Schlichtkrull et al., 2017](#schlichtkrull-et-al-2017)). * `RGAT`: Relational Graph Attention Networks ([Veličković et al., 2018](#veličković-et-al-2018)). -* `GNN-Edge-MLP`: Graph Neural Network with Edge MLPs - a variant of RGCN in which messages on edges are computed using full MLPs, not just a single layer. +* `GNN-Edge-MLP`: Graph Neural Network with Edge MLPs - a variant of RGCN in which messages on edges are computed using full MLPs, not just a single layer applied to the source state. * `RGDCN`: Relational Graph Dynamic Convolution Networks - a new variant of RGCN in which the weights of convolutional layers are dynamically computed. * `GNN-FiLM`: Graph Neural Networks with Feature-wise Linear Modulation - a new extension of RGCN with FiLM layers. @@ -130,13 +130,14 @@ by using `--data-path "SOME/OTHER/DIR"`. Running `python run_ppi_benchs.py ppi_results/` should yield results looking like this (on an NVidia V100): -| Model | Avg. MicroF1 | Avg. Time | -|--------------|-------------------|------------| -| GGNN | 0.990 (+/- 0.001) | 432.6 | -| RGCN | 0.989 (+/- 0.000) | 759.0 | -| GAT | 0.989 (+/- 0.001) | 782.3 | -| GNN-Edge-MLP | 0.992 (+/- 0.001) | 479.2 | -| GNN-FiLM | 0.992 (+/- 0.000) | 308.1 | +| Model | Avg. MicroF1 | Avg. Time | +|---------------|-------------------|------------| +| GGNN | 0.990 (+/- 0.001) | 432.6 | +| RGCN | 0.989 (+/- 0.000) | 759.0 | +| GAT | 0.989 (+/- 0.001) | 782.3 | +| GNN-Edge-MLP0 | 0.992 (+/- 0.000) | 556.9 | +| GNN-Edge-MLP1 | 0.992 (+/- 0.001) | 479.2 | +| GNN-FiLM | 0.992 (+/- 0.000) | 308.1 | ## QM9 The `QM9` task (implemented in `tasks/qm9_task.py`) handles the quantum chemistry diff --git a/models/ggnn_model.py b/models/ggnn_model.py index d07b15f..040f6d2 100644 --- a/models/ggnn_model.py +++ b/models/ggnn_model.py @@ -23,7 +23,7 @@ class GGNN_Model(Sparse_Graph_Model): return params @staticmethod - def name() -> str: + def name(params: Dict[str, Any]) -> str: return "GGNN" def __init__(self, params: Dict[str, Any], task: Sparse_Graph_Task, run_id: str, result_dir: str) -> None: diff --git a/models/gnn_edge_mlp_model.py b/models/gnn_edge_mlp_model.py index acf4297..36d078a 100644 --- a/models/gnn_edge_mlp_model.py +++ b/models/gnn_edge_mlp_model.py @@ -24,8 +24,8 @@ class GNN_Edge_MLP_Model(Sparse_Graph_Model): return params @staticmethod - def name() -> str: - return "GNN-Edge-MLP" + def name(params: Dict[str, Any]) -> str: + return "GNN-Edge-MLP%i" % (params['num_edge_hidden_layers']) def __init__(self, params: Dict[str, Any], task: Sparse_Graph_Task, run_id: str, result_dir: str) -> None: super().__init__(params, task, run_id, result_dir) diff --git a/models/gnn_film_model.py b/models/gnn_film_model.py index ddfc6c8..730d0d6 100644 --- a/models/gnn_film_model.py +++ b/models/gnn_film_model.py @@ -20,7 +20,7 @@ class GNN_FiLM_Model(Sparse_Graph_Model): return params @staticmethod - def name() -> str: + def name(params: Dict[str, Any]) -> str: return "GNN-FiLM" def __init__(self, params: Dict[str, Any], task: Sparse_Graph_Task, run_id: str, result_dir: str) -> None: diff --git a/models/rgat_model.py b/models/rgat_model.py index 0fc4c43..6665a7a 100644 --- a/models/rgat_model.py +++ b/models/rgat_model.py @@ -22,7 +22,7 @@ class RGAT_Model(Sparse_Graph_Model): return params @staticmethod - def name() -> str: + def name(params: Dict[str, Any]) -> str: return "RGAT" def __init__(self, params: Dict[str, Any], task: Sparse_Graph_Task, run_id: str, result_dir: str) -> None: diff --git a/models/rgcn_model.py b/models/rgcn_model.py index 3bee332..c3eca3d 100644 --- a/models/rgcn_model.py +++ b/models/rgcn_model.py @@ -22,7 +22,7 @@ class RGCN_Model(Sparse_Graph_Model): return params @staticmethod - def name() -> str: + def name(params: Dict[str, Any]) -> str: return "RGCN" def __init__(self, params: Dict[str, Any], task: Sparse_Graph_Task, run_id: str, result_dir: str) -> None: diff --git a/models/rgdcn_model.py b/models/rgdcn_model.py index 9c9a24c..a4fc0d2 100644 --- a/models/rgdcn_model.py +++ b/models/rgdcn_model.py @@ -24,7 +24,7 @@ class RGDCN_Model(Sparse_Graph_Model): return params @staticmethod - def name() -> str: + def name(params: Dict[str, Any]) -> str: return "RGDCN" def __init__(self, params: Dict[str, Any], task: Sparse_Graph_Task, run_id: str, result_dir: str) -> None: diff --git a/models/sparse_graph_model.py b/models/sparse_graph_model.py index 1c4e8c1..ac7c129 100644 --- a/models/sparse_graph_model.py +++ b/models/sparse_graph_model.py @@ -46,7 +46,7 @@ class Sparse_Graph_Model(ABC): @staticmethod @abstractmethod - def name() -> str: + def name(params: Dict[str, Any]) -> str: raise NotImplementedError() def __init__(self, @@ -96,7 +96,7 @@ class Sparse_Graph_Model(ABC): weights_to_save = self.sess.run(vars_to_retrieve) data_to_save = { - "model_class": self.name(), + "model_class": self.name(self.params), "task_class": self.task.name(), "model_params": self.params, "task_params": self.task.params, diff --git a/run_ppi_benchs.py b/run_ppi_benchs.py index eb784dc..c3ffd62 100644 --- a/run_ppi_benchs.py +++ b/run_ppi_benchs.py @@ -16,7 +16,7 @@ import numpy as np from docopt import docopt from dpu_utils.utils import run_and_debug -MODEL_TYPES = ["GGNN", "RGCN", "RGAT", "GNN-Edge-MLP", "GNN_FiLM"] +MODEL_TYPES = ["GGNN", "RGCN", "RGAT", "GNN-Edge-MLP0", "GNN-Edge-MLP1", "GNN_FiLM"] TEST_RES_RE = re.compile('^Metrics: Avg MicroF1: (0.\d+)') TIME_RE = re.compile('^Training took (\d+)s') diff --git a/tasks/default_hypers/PPI_GNN-Edge-MLP0.json b/tasks/default_hypers/PPI_GNN-Edge-MLP0.json new file mode 100644 index 0000000..988c4ed --- /dev/null +++ b/tasks/default_hypers/PPI_GNN-Edge-MLP0.json @@ -0,0 +1,7 @@ +{"task_params": {}, + "model_params": {"graph_num_layers": 5, + "hidden_size": 256, + "max_nodes_in_batch": 6000, + "graph_layer_input_dropout_keep_prob": 0.8 + } +} \ No newline at end of file diff --git a/tasks/default_hypers/PPI_GNN-Edge-MLP.json b/tasks/default_hypers/PPI_GNN-Edge-MLP1.json similarity index 100% rename from tasks/default_hypers/PPI_GNN-Edge-MLP.json rename to tasks/default_hypers/PPI_GNN-Edge-MLP1.json diff --git a/tasks/default_hypers/QM9_GNN-Edge-MLP.json b/tasks/default_hypers/QM9_GNN-Edge-MLP.json deleted file mode 100644 index c96d5a7..0000000 --- a/tasks/default_hypers/QM9_GNN-Edge-MLP.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "task_params": {}, - "model_params": { - "max_nodes_in_batch": 50000, - "num_edge_hidden_layers": 0 - } -} \ No newline at end of file diff --git a/tasks/default_hypers/QM9_GNN-Edge-MLP0.json b/tasks/default_hypers/QM9_GNN-Edge-MLP0.json new file mode 100644 index 0000000..c365708 --- /dev/null +++ b/tasks/default_hypers/QM9_GNN-Edge-MLP0.json @@ -0,0 +1,27 @@ +{ + "task_params": {}, + "model_params": { + "max_nodes_in_batch": 50000, + "graph_num_layers": 8, + "graph_num_timesteps_per_layer": 1, + "graph_layer_input_dropout_keep_prob": 0.9, + "graph_dense_between_every_num_gnn_layers": 32, + "graph_model_activation_function": "tanh", + "graph_residual_connection_every_num_layers": 2, + "graph_inter_layer_norm": true, + "max_epochs": 10000, + "patience": 25, + "optimizer": "RMSProp", + "learning_rate": 0.0005072060718321982, + "learning_rate_decay": 0.98, + "lr_for_num_graphs_per_batch": null, + "momentum": 0.85, + "clamp_gradient_norm": 1.0, + "hidden_size": 128, + "graph_activation_function": "relu", + "message_aggregation_function": "sum", + "graph_message_weights_dropout_ratio": 0.0, + "use_target_state_as_input": true, + "num_edge_hidden_layers": 0 + } +} \ No newline at end of file diff --git a/tasks/default_hypers/QM9_GNN-Edge-MLP1.json b/tasks/default_hypers/QM9_GNN-Edge-MLP1.json new file mode 100644 index 0000000..8fb73da --- /dev/null +++ b/tasks/default_hypers/QM9_GNN-Edge-MLP1.json @@ -0,0 +1,27 @@ +{ + "task_params": {}, + "model_params": { + "max_nodes_in_batch": 50000, + "graph_num_layers": 8, + "graph_num_timesteps_per_layer": 1, + "graph_layer_input_dropout_keep_prob": 0.9, + "graph_dense_between_every_num_gnn_layers": 32, + "graph_model_activation_function": "tanh", + "graph_residual_connection_every_num_layers": 2, + "graph_inter_layer_norm": false, + "max_epochs": 10000, + "patience": 25, + "optimizer": "Adam", + "learning_rate": 0.0006482335154980316, + "learning_rate_decay": 0.98, + "lr_for_num_graphs_per_batch": null, + "momentum": 0.85, + "clamp_gradient_norm": 1.0, + "hidden_size": 128, + "graph_activation_function": "gelu", + "message_aggregation_function": "sum", + "graph_message_weights_dropout_ratio": 0.0, + "use_target_state_as_input": true, + "num_edge_hidden_layers": 1 + } +} \ No newline at end of file diff --git a/train.py b/train.py index f9a3de2..375ba19 100644 --- a/train.py +++ b/train.py @@ -32,20 +32,21 @@ from test import test def run(args): azure_info_path = args.get('--azure-info', None) - model_cls = name_to_model_class(args['MODEL_NAME']) + model_cls, additional_model_params = name_to_model_class(args['MODEL_NAME']) task_cls, additional_task_params = name_to_task_class(args['TASK_NAME']) # Collect parameters from first the class defaults, potential task defaults, and then CLI: task_params = task_cls.default_params() task_params.update(additional_task_params) model_params = model_cls.default_params() + model_params.update(additional_model_params) # Load potential task-specific defaults: task_model_default_hypers_file = \ os.path.join(os.path.dirname(__file__), "tasks", "default_hypers", - "%s_%s.json" % (task_cls.name(), model_cls.name())) + "%s_%s.json" % (task_cls.name(), model_cls.name(model_params))) if os.path.exists(task_model_default_hypers_file): print("Loading task/model-specific default parameters from %s." % task_model_default_hypers_file) with open(task_model_default_hypers_file, "rt") as f: @@ -77,7 +78,7 @@ def run(args): for random_seed in random_seeds: model_params['random_seed'] = random_seed - run_id = "_".join([task_cls.name(), model_cls.name(), time.strftime("%Y-%m-%d-%H-%M-%S"), str(os.getpid())]) + run_id = "_".join([task_cls.name(), model_cls.name(model_params), time.strftime("%Y-%m-%d-%H-%M-%S"), str(os.getpid())]) model = model_cls(model_params, task, run_id, result_dir) model.log_line("Run %s starting." % run_id) diff --git a/utils/model_utils.py b/utils/model_utils.py index fb8fdf7..75cbdb8 100644 --- a/utils/model_utils.py +++ b/utils/model_utils.py @@ -29,34 +29,40 @@ def name_to_task_class(name: str) -> Tuple[Type[Sparse_Graph_Task], Dict[str, An raise ValueError("Unknown task type '%s'" % name) -def name_to_model_class(name: str) -> Type[Sparse_Graph_Model]: +def name_to_model_class(name: str) -> Tuple[Type[Sparse_Graph_Model], Dict[str, Any]]: name = name.lower() if name in ["ggnn", "ggnn_model"]: - return GGNN_Model + return GGNN_Model, {} if name in ["gnn_edge_mlp", "gnn-edge-mlp", "gnn_edge_mlp_model"]: - return GNN_Edge_MLP_Model + return GNN_Edge_MLP_Model, {} + if name in ["gnn_edge_mlp0", "gnn-edge-mlp0", "gnn_edge_mlp0_model"]: + return GNN_Edge_MLP_Model, {'num_edge_hidden_layers': 0} + if name in ["gnn_edge_mlp1", "gnn-edge-mlp1", "gnn_edge_mlp1_model"]: + return GNN_Edge_MLP_Model, {'num_edge_hidden_layers': 1} + if name in ["gnn_edge_mlp", "gnn-edge-mlp"]: + return GNN_Edge_MLP_Model, {} if name in ["gnn_film", "gnn-film", "gnn_film_model"]: - return GNN_FiLM_Model + return GNN_FiLM_Model, {} if name in ["rgat", "rgat_model"]: - return RGAT_Model + return RGAT_Model, {} if name in ["rgcn", "rgcn_model"]: - return RGCN_Model + return RGCN_Model, {} if name in ["rgdcn", "rgdcn_model"]: - return RGDCN_Model + return RGDCN_Model, {} raise ValueError("Unknown model type '%s'" % name) -def restore(saved_model_path: str, result_dir: str, run_id: str = None) -> None: +def restore(saved_model_path: str, result_dir: str, run_id: str = None) -> Sparse_Graph_Model: print("Loading model from file %s." % saved_model_path) with open(saved_model_path, 'rb') as in_file: data_to_load = pickle.load(in_file) - model_cls = name_to_model_class(data_to_load['model_class']) + model_cls, _ = name_to_model_class(data_to_load['model_class']) task_cls, additional_task_params = name_to_task_class(data_to_load['task_class']) if run_id is None: - run_id = "_".join([task_cls.name(), model_cls.name(), time.strftime("%Y-%m-%d-%H-%M-%S"), str(os.getpid())]) + run_id = "_".join([task_cls.name(), model_cls.name(data_to_load['model_params']), time.strftime("%Y-%m-%d-%H-%M-%S"), str(os.getpid())]) task = task_cls(data_to_load['task_params']) task.restore_from_metadata(data_to_load['task_metadata'])