Heterogeneous Graph Support + GraphGym

.github/workflows/testing.yml

@@ -55,7 +55,9 @@ jobs:
           pip install pandas
           pip install h5py
           pip install numba
+          pip install tabulate
           pip install git+https://github.com/the-virtual-brain/tvb-geodesic.git
+          pip install pytorch-memlab
 
       - name: Install main package
         run: |

.gitignore

@@ -8,6 +8,7 @@ runs/
 .eggs/
 lightning_logs/
 *.egg-info/
+.ipynb_checkpoints
 .coverage
 .coverage.*
 coverage.xml

LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2020 Matthias Fey <[email protected]>
+Copyright (c) 2021 Matthias Fey, Jiaxuan You <[email protected], [email protected]>
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

benchmark/kernel/README.md

@@ -21,3 +21,9 @@ Run (or modify) the whole test suite via
 ```
 $ python main.py
 ```
+
+For more comprehensive time-measurement and memory usage information, you may use 
+
+```
+$ python main_performance.py
+```

benchmark/kernel/main_performance.py

@@ -0,0 +1,79 @@
+import argparse
+from itertools import product
+
+import torch
+
+from torch_geometric import seed_everything
+from torch_geometric.loader import DataLoader
+from torch_geometric.profile import profileit, timeit, get_stats_summary
+
+from datasets import get_dataset
+from train_eval import train, eval_acc
+
+from gcn import GCN
+from graph_sage import GraphSAGE
+from gin import GIN
+
+seed_everything(0)
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--epochs', type=int, default=100)
+parser.add_argument('--batch_size', type=int, default=128)
+parser.add_argument('--lr', type=float, default=0.01)
+parser.add_argument('--warmup_profile', type=int, default=1,
+                    help='Skip the first few runs')
+parser.add_argument('--goal_accuracy', type=int, default=1,
+                    help='The goal test accuracy')
+args = parser.parse_args()
+
+layers = [1, 2, 3]
+hiddens = [16, 32]
+
+datasets = ['MUTAG', 'PROTEINS', 'IMDB-BINARY', 'REDDIT-BINARY']
+
+nets = [
+    GCN,
+    GraphSAGE,
+    GIN,
+]
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Decorate train and eval functions:
+train = profileit(print_layer_stats=False)(train)
+eval_acc = timeit()(eval_acc)
+
+for dataset_name, Net in product(datasets, nets):
+    dataset = get_dataset(dataset_name, sparse=True)
+    num_train = int(len(dataset) * 0.8)
+    num_val = int(len(dataset) * 0.1)
+
+    train_dataset = dataset[:num_train]
+    val_dataset = dataset[num_train:num_train + num_val]
+    test_dataset = dataset[num_train + num_val:]
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
+                              shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
+                            shuffle=False)
+    test_loader = DataLoader(test_dataset, batch_size=args.batch_size,
+                             shuffle=False)
+
+    for num_layers, hidden in product(layers, hiddens):
+        print('-----\n{} - {} - {} - {}'.format(dataset_name, Net.__name__,
+                                                num_layers, hidden))
+
+        model = Net(dataset, num_layers, hidden).to(device)
+        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+        stats_list = []
+        for epoch in range(1, args.epochs + 1):
+            loss, stats = train(model, optimizer, train_loader)
+            val_acc, val_time = eval_acc(model, val_loader)
+            test_acc, test_time = eval_acc(model, test_loader)
+
+            if epoch >= args.warmup_profile:
+                stats_list.append(stats)
+
+        stats_summary = get_stats_summary(stats_list)
+        print(stats_summary)

benchmark/kernel/train_eval.py

@@ -5,7 +5,7 @@
 from torch import tensor
 from torch.optim import Adam
 from sklearn.model_selection import StratifiedKFold
-from torch_geometric.data import DataLoader, DenseDataLoader as DenseLoader
+from torch_geometric.loader import DataLoader, DenseDataLoader as DenseLoader
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 

benchmark/points/train_eval.py

@@ -3,7 +3,7 @@
 import torch
 import torch.nn.functional as F
 from torch.optim import Adam
-from torch_geometric.data import DataLoader
+from torch_geometric.loader import DataLoader
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 

docs/source/_figures/edge_conv.tex

@@ -0,0 +1,12 @@
+\documentclass{standalone}
+
+\usepackage{bm}
+\usepackage{tikz}
+
+\begin{document}
+
+\begin{tikzpicture}
+  \node at (0,0) {$\bm{x}_i^{\prime} = \max_{j \in \mathcal{N}(i)} \textsc{MLP}_{\hspace{-1pt}\theta} \big([ \bm{x}_i, \bm{x}_j - \bm{x}_i ]\big)$};
+\end{tikzpicture}
+
+\end{document}