Refactoring

Author: urialon

README.md

@@ -3,33 +3,35 @@
 This is the official implementation of the paper: [On the Bottleneck of Graph Neural Networks and its Practical Implications](https://openreview.net/pdf?id=i80OPhOCVH2) (ICLR'2021). 
 
 By [Uri Alon](http://urialon.cswp.cs.technion.ac.il/) and [Eran Yahav](http://www.cs.technion.ac.il/~yahave/).
-See also the [[video]](https://youtu.be/vrLsEwzZTCQ) and the [[slides]](https://urialon.cswp.cs.technion.ac.il/wp-content/uploads/sites/83/2020/07/bottleneck_slides.pdf). 
+See also the [[video]](https://youtu.be/vrLsEwzZTCQ), [[poster]](https://urialon.cswp.cs.technion.ac.il/wp-content/uploads/sites/83/2021/03/bottleneck_poster.pdf) and [[slides]](https://urialon.cswp.cs.technion.ac.il/wp-content/uploads/sites/83/2020/07/bottleneck_slides.pdf). 
 
 this repository is divided into three sub-projects:
 
 1. The subdirectory `tf-gnn-samples` is a clone of 
 [https://github.com/microsoft/tf-gnn-samples](https://github.com/microsoft/tf-gnn-samples) by Brockschmidt (ICML'2020). 
+This project can be used to reproduce the QM9 and VarMisuse experiments of Section 4.2 and 4.2 in the paper.
 This sub-project depends on TensorFlow 1.13.
 The instructions for our clone are the same as their original code, except that reproducing our experiments 
 (the QM9 dataset and VarMisuse) can be done by running the
 script `tf-gnn-samples/run_qm9_benchs_fa.py` or `tf-gnn-samples/run_varmisuse_benchs_fa.py` instead of their original scripts.
 For additional dependencies and instructions, see their original README: 
 [https://github.com/microsoft/tf-gnn-samples/blob/master/README.md](https://github.com/microsoft/tf-gnn-samples/blob/master/README.md).
 The main modification that we performed is using a Fully-Adjacent layer as the last 
-GNN layer and we describe in our paper (Section 4). 
+GNN layer and we describe in our paper. 
 2. The subdirectory `gnn-comparison` is a clone of [https://github.com/diningphil/gnn-comparison](https://github.com/diningphil/gnn-comparison)
-by Errica et al. (ICLR'2020). This sub-project depends on PyTorch 1.4 and Pytorch-Geometric.
+by Errica et al. (ICLR'2020). 
+This project can be used to reproduce the biological experiments (Section 4.3, the ENZYMES and NCI1 datasets).
+This sub-project depends on PyTorch 1.4 and Pytorch-Geometric.
 For additional dependencies and instructions, see their original README: 
 [https://github.com/diningphil/gnn-comparison/blob/master/README.md](https://github.com/diningphil/gnn-comparison/blob/master/README.md).
-The instructions for our clone are the same, except that we added an additional flag to every `config_*.yml` file, called `last_layer_fully_adjacent`, 
-which is set to `True` by default, and reproduces our experiments (Section 4.3, the ENZYMES and NCI1 datasets).
+The instructions for our clone are the same, except that we added an additional flag to every `config_*.yml` file, called `last_layer_fa`, 
+which is set to `True` by default, and reproduces our experiments.
 The main modification that we performed is using a Fully-Adjacent layer as the last 
 GNN layer.
 3. The main directory (in which this file resides) can be used to reproduce the experiments of 
 Section 4.1 in the paper, for the "Tree-NeighborsMatch" problem. The rest of this README file includes the 
 instructions for this main directory. 
 This repository can be used to reproduce the experiments of 
-Section 4.1 in the paper, for the "Tree-NeighborsMatch" problem. 
 
 This project was designed to be useful in experimenting with new GNN architectures and new solutions for the over-squashing problem. 
 

common.py

@@ -7,13 +7,7 @@
 
 
 class Task(Enum):
-    PARITY = auto()
-    LEAF_PARITY = auto()
-    NOISY_LEAF_PARITY = auto()
-    NOISY_SUM = auto()
-    DICTIONARY = auto()
-    NOISY_DICTIONARY = auto()
-    DUMMY = auto()
+    NEIGHBORS_MATCH = auto()
 
     @staticmethod
     def from_string(s):
@@ -23,7 +17,7 @@ def from_string(s):
             raise ValueError()
 
     def get_dataset(self, depth, train_fraction):
-        if self is Task.DICTIONARY:
+        if self is Task.NEIGHBORS_MATCH:
             dataset = DictionaryLookupDataset(depth)
         else:
             dataset = None
@@ -55,7 +49,7 @@ def get_layer(self, in_dim, out_dim):
             return GINConv(nn.Sequential(nn.Linear(in_dim, out_dim), nn.BatchNorm1d(out_dim), nn.ReLU(),
                                          nn.Linear(out_dim, out_dim), nn.BatchNorm1d(out_dim), nn.ReLU()))
         elif self is GNN_TYPE.GAT:
-            # 4-heads, although the paper by Velickovic et al. had used 6-8 heads. 
+            # 4-heads, although the paper by Velickovic et al. had used 6-8 heads.
             # The output will be the concatenation of the heads, yielding a vector of size out_dim
             num_heads = 4
             return GATConv(in_dim, out_dim // num_heads, heads=num_heads)

main.py

@@ -6,7 +6,7 @@
 
 if __name__ == '__main__':
     parser = ArgumentParser()
-    parser.add_argument("--task", dest="task", default=Task.PARITY, type=Task.from_string, choices=list(Task),
+    parser.add_argument("--task", dest="task", default=Task.NEIGHBORS_MATCH, type=Task.from_string, choices=list(Task),
                         required=False)
     parser.add_argument("--type", dest="type", default=GNN_TYPE.GCN, type=GNN_TYPE.from_string, choices=list(GNN_TYPE),
                         required=False)
@@ -18,7 +18,7 @@
     parser.add_argument("--eval_every", dest="eval_every", default=100, type=int, required=False)
     parser.add_argument("--batch_size", dest="batch_size", default=1024, type=int, required=False)
     parser.add_argument("--accum_grad", dest="accum_grad", default=1, type=int, required=False)
-    parser.add_argument("--stop", dest="stop", default=STOP.TRAIN, type=STOP.from_string, choices=list(STOP), 
+    parser.add_argument("--stop", dest="stop", default=STOP.TRAIN, type=STOP.from_string, choices=list(STOP),
                         required=False)
     parser.add_argument("--patience", dest="patience", default=20, type=int, required=False)
     parser.add_argument("--loader_workers", dest="loader_workers", default=0, type=int, required=False)
@@ -33,7 +33,7 @@
 
 
 def get_fake_args(
-        task=Task.PARITY,
+        task=Task.NEIGHBORS_MATCH,
         type=GNN_TYPE.GCN,
         dim=32,
         depth=3,

run-gat-2-8.py

@@ -24,7 +24,7 @@ def __init__(self, train_acc, test_acc, epoch):
 
 if __name__ == '__main__':
 
-    task = Task.DICTIONARY
+    task = Task.NEIGHBORS_MATCH
     gnn_type = GNN_TYPE.GAT
     stopping_criterion = STOP.TRAIN
     min_depth = 2

run-gcn-2-8.py

@@ -24,7 +24,7 @@ def __init__(self, train_acc, test_acc, epoch):
 
 if __name__ == '__main__':
 
-    task = Task.DICTIONARY
+    task = Task.NEIGHBORS_MATCH
     gnn_type = GNN_TYPE.GCN
     stopping_criterion = STOP.TRAIN
     min_depth = 2

run-ggnn-2-8.py

@@ -24,7 +24,7 @@ def __init__(self, train_acc, test_acc, epoch):
 
 if __name__ == '__main__':
 
-    task = Task.DICTIONARY
+    task = Task.NEIGHBORS_MATCH
     gnn_type = GNN_TYPE.GGNN
     stopping_criterion = STOP.TRAIN
     min_depth = 2

run-gin-2-8.py

@@ -24,7 +24,7 @@ def __init__(self, train_acc, test_acc, epoch):
 
 if __name__ == '__main__':
 
-    task = Task.DICTIONARY
+    task = Task.NEIGHBORS_MATCH
     gnn_type = GNN_TYPE.GIN
     stopping_criterion = STOP.TRAIN
     min_depth = 2