initial push alchemy

Author: cptq

Alchemy/README.md

@@ -0,0 +1,21 @@
+## Graph Regression Experiments on Alchemy
+
+### Usage
+
+To run the experiments for SignNet, use `python main_alchemy.py`.
+
+### Implementation
+
+Our SignNet model is implemented in PyTorch Geometric in the `sign_net` folder.
+
+### Setup
+
+Requirements are in `setup.sh`. Simply running `bash setup.sh` will usually make a conda environment called `torch-1-9` that works for these experiments, which you can then activate with `conda activate torch-1-9`.
+
+You may have to edit the `CUDA` variable in `setup.sh` depending on the CUDA version of your GPUs. We use PyTorch 1.9 and PyTorch Geometric 2.0.1.
+
+### Attribution
+
+We built off of the SpeqNets repo by Christopher Morris et al. (no license) [[link](https://github.com/chrsmrrs/SpeqNets/blob/master/neural_graph/main_1_alchemy_10K.py)].
+
+The Alchemy dataset is from "Alchemy: A Quantum Chemistry Dataset for Benchmarking AI Models" Chen et al. 2019  [[arXiv link](https://arxiv.org/abs/1906.09427)].

Alchemy/baseline_gin.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+from torch.nn import Sequential, Linear, ReLU
+import torch.nn.functional as F
+
+from torch_geometric.nn import MessagePassing, Set2Set
+
+
+class GINConv(MessagePassing):
+    def __init__(self, emb_dim, dim1, dim2):
+        super(GINConv, self).__init__(aggr="add")
+
+        self.bond_encoder = Sequential(Linear(emb_dim, dim1), ReLU(), Linear(dim1, dim1))
+        self.mlp = Sequential(Linear(dim1, dim1), ReLU(), Linear(dim1, dim2))
+        self.eps = nn.Parameter(torch.Tensor([0]))
+
+    def forward(self, x, edge_index, edge_attr):
+        edge_embedding = self.bond_encoder(edge_attr)
+        out = self.mlp((1 + self.eps) * x + self.propagate(edge_index, x=x, edge_attr=edge_embedding))
+        return out
+
+    def message(self, x_j, edge_attr):
+        return F.relu(x_j + edge_attr)
+
+    def update(self, aggr_out):
+        return aggr_out
+
+
+class NetGINE(nn.Module):
+    def __init__(self, dim):
+        super(NetGINE, self).__init__()
+
+        num_features = 6
+        dim = dim
+
+        self.conv1 = GINConv(4, num_features, dim)
+        self.conv2 = GINConv(4, dim, dim)
+        self.conv3 = GINConv(4, dim, dim)
+        self.conv4 = GINConv(4, dim, dim)
+        self.conv5 = GINConv(4, dim, dim)
+        self.conv6 = GINConv(4, dim, dim)
+
+        self.set2set = Set2Set(1 * dim, processing_steps=6)
+
+        self.fc1 = Linear(2 * dim, dim)
+        self.fc4 = Linear(dim, 12)
+
+    def forward(self, data):
+        x = data.x
+
+        x_1 = F.relu(self.conv1(x, data.edge_index, data.edge_attr))
+        x_2 = F.relu(self.conv2(x_1, data.edge_index, data.edge_attr))
+        x_3 = F.relu(self.conv3(x_2, data.edge_index, data.edge_attr))
+        x_4 = F.relu(self.conv4(x_3, data.edge_index, data.edge_attr))
+        x_5 = F.relu(self.conv5(x_4, data.edge_index, data.edge_attr))
+        x_6 = F.relu(self.conv6(x_5, data.edge_index, data.edge_attr))
+        x = x_6
+        x = self.set2set(x, data.batch)
+        x = F.relu(self.fc1(x))
+        x = self.fc4(x)
+        return x

Alchemy/setup.sh

@@ -0,0 +1,19 @@
+# modified from setup in https://github.com/LingxiaoShawn/GNNAsKernel
+ENV=torch-1-9
+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
+
+# 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 matplotlib

Alchemy/sign_net/__init__.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_scatter import scatter
+import sign_net.model_utils.pyg_gnn_wrapper as gnn_wrapper 
+from sign_net.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, res=True):
+        super().__init__()
+        self.input_encoder = DiscreteEncoder(nhid) if nfeat_node is None else MLP(nfeat_node, nhid, 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)
+
+
+        self.pooling = pooling
+        self.dropout = dropout
+        self.res = res
+
+    def reset_parameters(self):
+        self.input_encoder.reset_parameters()
+        self.output_encoder.reset_parameters()
+        #self.size_embedder.reset_parameters()
+        self.linear.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())
+
+        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 = 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')
+
+        x = self.output_encoder(x)
+        return x
+

Alchemy/sign_net/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=6): #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 
+
+

Alchemy/sign_net/model_utils/__init__.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+import torch_geometric.nn as gnn
+import torch.nn.functional as F
+from sign_net.model_utils.elements import Identity
+
+class MaskedBN(nn.Module):
+    def __init__(self, num_features):
+        super().__init__()
+        self.bn = nn.BatchNorm1d(num_features)
+    def reset_parameters(self):
+        self.bn.reset_parameters() 
+    def forward(self, x, mask=None):
+        ### apply BN to the last dim
+        #    x: n x k x d
+        # mask: n x k  
+        if mask is None:
+            return self.bn(x.transpose(1,2)).transpose(1,2)
+        x[mask] = self.bn(x[mask])
+        return x
+
+class MaskedLN(nn.Module):
+    def __init__(self, num_features):
+        super().__init__()
+        self.ln = nn.LayerNorm(num_features, eps=1e-6)
+    def reset_parameters(self):
+        self.ln.reset_parameters() 
+    def forward(self, x, mask=None):
+        if mask is None:
+            return self.ln(x)
+        x[mask] = self.ln(x[mask])
+        return x
+
+class MaskedMLP(nn.Module):
+    def __init__(self, nin, nout, nlayer=2, with_final_activation=True, with_norm=True, 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([MaskedBN(n_hid if i<nlayer-1 else nout) 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, mask=None):
+        # x: n x k x d
+        previous_x = x
+        for i, (layer, norm) in enumerate(zip(self.layers, self.norms)):
+            x = layer(x)
+            if mask is not None:
+                x[~mask] = 0
+            if i < self.nlayer-1 or self.with_final_activation:
+                x = norm(x, mask)
+                x = F.relu(x)  
+        return x    
+
+class MaskedGINConv(nn.Module):
+    def __init__(self, nin, nout, bias=True, nhid=None):
+        super().__init__()
+        self.nn = MaskedMLP(nin, nout, 2, False, bias=bias, nhid=nhid)
+        self.layer = gnn.GINConv(Identity(), train_eps=True)
+    def reset_parameters(self):
+        self.nn.reset_parameters()
+        self.layer.reset_parameters()
+    def forward(self, x, edge_index, edge_attr, mask=None):
+        x = self.layer(x, edge_index)
+        if mask is not None:
+            if x[~mask].numel() == 0:
+                print('~mask numel = 0!!')
+                print('x shape', x.shape)
+                print('mask shape', mask.shape)
+            #assert x[~mask].max() == 0 
+        x = self.nn(x, mask)
+        # assert x[~mask].max() == 0 
+        return x
+
+
+class MaskedGINEConv(nn.Module):
+    def __init__(self, nin, nout, bias=True, nhid=None):
+        super().__init__()
+        self.nn = MaskedMLP(nin, nout, 2, False, bias=bias, nhid=nhid)
+        self.layer = gnn.GINEConv(Identity(), train_eps=True)
+    def reset_parameters(self):
+        self.nn.reset_parameters()
+        self.layer.reset_parameters()
+    def forward(self, x, edge_index, edge_attr, mask=None):
+        assert x[~mask].max() == 0
+        x = self.layer(x, edge_index, edge_attr)
+        if mask is not None:
+            x[~mask] = 0 
+        x = self.nn(x, mask)
+        # assert x[~mask].max() == 0 
+        return x