add gine signnet zinc

Author: cptq

GINESignNetPyG/README.md

@@ -0,0 +1,43 @@
+## SignNet and GINE implentations in PyTorch Geometric
+
+For reproduction of our results on ZINC in the paper.
+
+This is approximately the same as our code used for the Alchemy experiments.
+
+## Setup 
+
+```
+# params
+# 10/6/2021, newest packages. 
+ENV=pyg
+CUDA=11.1
+TORCH=1.9.1
+PYG=2.0.1
+
+# create env 
+conda create --name $ENV python=3.9 -y
+conda activate $ENV
+
+# install pytorch 
+conda install pytorch=$TORCH torchvision torchaudio cudatoolkit=$cuda -c pytorch -c nvidia -y
+
+# install pyg2.0
+conda install pyg=$PYG -c pyg -c conda-forge -y
+
+# install ogb 
+pip install ogb
+
+# install rdkit
+conda install -c conda-forge rdkit -y
+
+# update yacs and tensorboard
+pip install yacs==0.1.8 --force  # PyG currently use 0.1.6 which doesn't support None argument. 
+pip install tensorboard
+pip install matplotlib
+
+```
+
+## Run
+```
+python -m train.zinc model.gnn_type SignNet
+```

GINESignNetPyG/core/__init__.py

@@ -0,0 +1,97 @@
+from yacs.config import CfgNode as CN
+
+def set_cfg(cfg):
+
+    # ------------------------------------------------------------------------ #
+    # Basic options
+    # ------------------------------------------------------------------------ #
+    # Dataset name
+    cfg.dataset = 'ZINC'
+    # Additional num of worker for data loading
+    cfg.num_workers = 8
+    # Cuda device number, used for machine with multiple gpus
+    cfg.device = 0 
+    # Additional string add to logging 
+    cfg.handtune = ''
+    # Whether fix the running seed to remove randomness
+    cfg.seed = None
+    # Whether downsampling the dataset, used for large dataset for faster tuning
+    cfg.downsample = False 
+    # version 
+    cfg.version = 'final'
+    # task, for simulation datasets
+    cfg.task = -1
+
+    # ------------------------------------------------------------------------ #
+    # Training options
+    # ------------------------------------------------------------------------ #
+    cfg.train = CN()
+    # Total graph mini-batch size
+    cfg.train.batch_size = 128
+    # Maximal number of epochs
+    cfg.train.epochs = 100
+    # Number of runs with random init 
+    cfg.train.runs = 3
+    # Base learning rate
+    cfg.train.lr = 0.001
+    # number of steps before reduce learning rate
+    cfg.train.lr_patience = 50
+    # learning rate decay factor
+    cfg.train.lr_decay = 0.5
+    # L2 regularization, weight decay
+    cfg.train.wd = 0.
+    # Dropout rate
+    cfg.train.dropout = 0.
+    
+    # ------------------------------------------------------------------------ #
+    # Model options
+    # ------------------------------------------------------------------------ #
+    cfg.model = CN()
+    # GNN type used, see core.model_utils.pyg_gnn_wrapper for all options
+    cfg.model.gnn_type = 'GINEConv' # change to list later
+    # Hidden size of the model
+    cfg.model.hidden_size = 128
+    # Number of gnn layers (doesn't include #MLPs)
+    cfg.model.num_layers = 4
+    # Number of signnet layers
+    cfg.model.num_layers_sign = 4
+    # Pooling type for generaating graph/subgraph embedding from node embeddings 
+    cfg.model.pool = 'add'
+
+    return cfg
+    
+import os 
+import argparse
+# Principle means that if an option is defined in a YACS config object, 
+# then your program should set that configuration option using cfg.merge_from_list(opts) and not by defining, 
+# for example, --train-scales as a command line argument that is then used to set cfg.TRAIN.SCALES.
+
+def update_cfg(cfg, args_str=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', default="", metavar="FILE", help="Path to config file")
+    # opts arg needs to match set_cfg
+    parser.add_argument("opts", default=[], nargs=argparse.REMAINDER, 
+                         help="Modify config options using the command-line")
+
+    if isinstance(args_str, str):
+        # parse from a string
+        args = parser.parse_args(args_str.split())
+    else:
+        # parse from command line
+        args = parser.parse_args()
+    # Clone the original cfg 
+    cfg = cfg.clone()
+    
+    # Update from config file
+    if os.path.isfile(args.config):
+        cfg.merge_from_file(args.config)
+
+    # Update from command line 
+    cfg.merge_from_list(args.opts)
+       
+    return cfg
+
+"""
+    Global variable
+"""
+cfg = set_cfg(CN())

GINESignNetPyG/core/config.py

@@ -0,0 +1,43 @@
+# Create a simple logger that can log training curves and final performance 
+from torch.utils.tensorboard import SummaryWriter # tensorboard
+import logging, os, sys, shutil
+import datetime
+
+def config_logger(cfg, OUT_PATH="results/", time=False):
+    # time option is used for debugging different model architecture. 
+    data_name = cfg.dataset 
+    if cfg.handtune:
+        data_name += f'-{cfg.handtune}'
+    # generate config_string
+    os.makedirs(os.path.join(OUT_PATH, cfg.version), exist_ok=True)
+    config_string = f'T[{cfg.task}] GNN[{cfg.model.gnn_type}] L[{cfg.model.num_layers_sign}-{cfg.model.num_layers}]  '\
+                    f'H[{cfg.model.hidden_size}] Pool[{cfg.model.pool}] '\
+                    f'Reg[{cfg.train.dropout}-{cfg.train.wd}] Seed[{cfg.seed}] GPU[{cfg.device}]'
+    
+    # setup tensorboard writer
+    writer_folder = os.path.join(OUT_PATH, cfg.version, data_name, config_string)
+    if time:
+        writer_folder = os.path.join(writer_folder, datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
+    if os.path.isdir(writer_folder): shutil.rmtree(writer_folder) # reset the folder, can also not reset
+    writer = SummaryWriter(writer_folder)
+
+    # setup logging
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    logger_filer = os.path.join(OUT_PATH, cfg.version, data_name, 'summary.log')
+    fh = logging.FileHandler(logger_filer)
+    fh.setLevel(logging.INFO)
+    fh.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(fh)
+
+    # redirect stdout print, better for large scale experiments
+    os.makedirs(os.path.join('logs', data_name), exist_ok=True)
+    # sys.stdout = open(f'logs/{data_name}/{config_string}.txt', 'w')
+
+    # log configuration 
+    print("-"*50)
+    print(cfg)
+    print("-"*50)
+    print('Time:', datetime.datetime.now().strftime("%Y/%m/%d - %H:%M"))
+    print(config_string)
+    return writer, logger, config_string

GINESignNetPyG/core/log.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_scatter import scatter
+import core.model_utils.pyg_gnn_wrapper as gnn_wrapper 
+from core.model_utils.elements import MLP, DiscreteEncoder, Identity, BN
+from torch_geometric.nn.inits import reset
+
+class GNN(nn.Module):
+    # this version use nin as hidden instead of nout, resulting a larger model
+    def __init__(self, nfeat_node, nfeat_edge, nhid, nout, nlayer, gnn_type, dropout=0, pooling='add', bn=BN, dos_bins=0, res=True):
+        super().__init__()
+        self.input_encoder = DiscreteEncoder(nhid-dos_bins) if nfeat_node is None else MLP(nfeat_node, nhid-dos_bins, 1)
+        self.edge_encoders = nn.ModuleList([DiscreteEncoder(nhid) if nfeat_edge is None else MLP(nfeat_edge, nhid, 1) for _ in range(nlayer)])
+        self.convs = nn.ModuleList([getattr(gnn_wrapper, gnn_type)(nhid, nhid, bias=not bn) for _ in range(nlayer)]) # set bias=False for BN
+        self.norms = nn.ModuleList([nn.BatchNorm1d(nhid) if bn else Identity() for _ in range(nlayer)])
+        self.output_encoder = MLP(nhid, nout, nlayer=2, with_final_activation=False, with_norm=False if pooling=='mean' else True)
+        self.size_embedder = nn.Embedding(200, nhid) 
+        self.linear = nn.Linear(2*nhid, nhid)
+
+        if dos_bins > 0:
+            self.ldos_encoder = MLP(dos_bins, nhid, nlayer=2, with_final_activation=True, with_norm=True)
+            self.dos_encoder = MLP(dos_bins, nhid, nlayer=2, with_final_activation=False, with_norm=True)
+
+        self.pooling = pooling
+        self.dropout = dropout
+        self.res = res
+        # for additional feature from (L)DOS
+        self.dos_bins = dos_bins
+
+    def reset_parameters(self):
+        self.input_encoder.reset_parameters()
+        self.output_encoder.reset_parameters()
+        self.size_embedder.reset_parameters()
+        self.linear.reset_parameters()
+        if self.dos_bins > 0:
+            self.dos_encoder.reset_parameters()
+            self.ldos_encoder.reset_parameters()
+        for edge_encoder, conv, norm in zip(self.edge_encoders, self.convs, self.norms):
+            edge_encoder.reset_parameters()
+            conv.reset_parameters()
+            norm.reset_parameters()
+     
+    def forward(self, data, additional_x=None):
+        x = self.input_encoder(data.x.squeeze())
+
+        # for PDOS 
+        if self.dos_bins > 0:
+            x = torch.cat([x, data.pdos], dim=-1)
+            # x += self.ldos_encoder(data.pdos)
+
+        if additional_x is not None:
+            x = self.linear(torch.cat([x, additional_x], dim=-1))
+
+        ori_edge_attr = data.edge_attr 
+        if ori_edge_attr is None:
+            ori_edge_attr = data.edge_index.new_zeros(data.edge_index.size(-1))
+
+        previous_x = x
+        for edge_encoder, layer, norm in zip(self.edge_encoders, self.convs, self.norms):
+            edge_attr = edge_encoder(ori_edge_attr) 
+            x = layer(x, data.edge_index, edge_attr)
+            x = norm(x)
+            x = F.relu(x)
+            x = F.dropout(x, self.dropout, training=self.training)
+            if self.res:
+                x += previous_x 
+                previous_x = x
+
+        if self.pooling == 'mean':
+            graph_size = scatter(torch.ones_like(x[:,0], dtype=torch.int64), data.batch, dim=0, reduce='add')
+            x = scatter(x, data.batch, dim=0, reduce='mean') + self.size_embedder(graph_size)
+        else:
+            x = scatter(x, data.batch, dim=0, reduce='add')
+
+        if self.dos_bins > 0:
+            x = x + self.dos_encoder(data.dos)
+        x = self.output_encoder(x)
+        return x
+

GINESignNetPyG/core/model.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import global_add_pool
+
+BN = True
+# BN = False
+RUNNING_STAT = True
+
+
+class Identity(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+    def reset_parameters(self):
+        pass
+
+class DiscreteEncoder(nn.Module):
+    def __init__(self, hidden_channels, max_num_features=10, max_num_values=500): #10
+        super().__init__()
+        self.embeddings = nn.ModuleList([nn.Embedding(max_num_values, hidden_channels) 
+                    for i in range(max_num_features)])
+
+    def reset_parameters(self):
+        for embedding in self.embeddings:
+            embedding.reset_parameters()
+            
+    def forward(self, x):
+        if x.dim() == 1:
+            x = x.unsqueeze(1)
+        out = 0
+        for i in range(x.size(1)):
+            out += self.embeddings[i](x[:, i])
+        return out
+
+class MLP(nn.Module):
+    def __init__(self, nin, nout, nlayer=2, with_final_activation=True, with_norm=BN, bias=True, nhid=None):
+        super().__init__()
+        n_hid = nin if nhid is None else nhid
+        self.layers = nn.ModuleList([nn.Linear(nin if i==0 else n_hid, 
+                                     n_hid if i<nlayer-1 else nout, 
+                                     bias=True if (i==nlayer-1 and not with_final_activation and bias) # TODO: revise later
+                                        or (not with_norm) else False) # set bias=False for BN
+                                     for i in range(nlayer)])
+        self.norms = nn.ModuleList([nn.BatchNorm1d(n_hid if i<nlayer-1 else nout,track_running_stats=RUNNING_STAT) if with_norm else Identity()
+                                     for i in range(nlayer)])
+        self.nlayer = nlayer
+        self.with_final_activation = with_final_activation
+        self.residual = (nin==nout) ## TODO: test whether need this
+
+    def reset_parameters(self):
+        for layer, norm in zip(self.layers, self.norms):
+            layer.reset_parameters()
+            norm.reset_parameters()
+
+    def forward(self, x):
+        previous_x = x
+        for i, (layer, norm) in enumerate(zip(self.layers, self.norms)):
+            x = layer(x)
+            if i < self.nlayer-1 or self.with_final_activation:
+                x = norm(x)
+                x = F.relu(x)  
+
+        # if self.residual:
+        #     x = x + previous_x  
+        return x 
+
+