[Benchmark] Add inference test for current benchmarking and record pe…

…rformance (#4892)

* [Benchmark] Add inference and profile in citation

* Print end-to-end time of inference

* Print end-to-end time of one epoch

* Add inference.sh

* Add inference and profile for to_hetero_mag

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add inference for pna

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add inference for benchmark/points/edge_cnn

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix error

* Update scripts

* Add profile test to increase code coverage

* Update script of points benchmark

* Update script for missing rename

* Update scripts according to the comments

* Add CPU test for profile

* update

* update

* update

* update

* reset

* update

* changelog

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.0.5] - 2022-MM-DD
 ### Added
+- Added inference benchmarks ([#4892](https://github.com/pyg-team/pytorch_geometric/pull/4892))
 - Added PyTorch 1.12 support  ([#4975](https://github.com/pyg-team/pytorch_geometric/pull/4975))
 - Added `unbatch_edge_index` functionality for splitting an `edge_index` tensor according to a `batch` vector ([#4903](https://github.com/pyg-team/pytorch_geometric/pull/4903))
 - Added node-wise normalization mode in `LayerNorm` ([#4944](https://github.com/pyg-team/pytorch_geometric/pull/4944))

benchmark/citation/appnp.py

@@ -6,6 +6,7 @@
 from torch.nn import Linear
 
 from torch_geometric.nn import APPNP
+from torch_geometric.profile import rename_profile_file
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, required=True)
@@ -20,6 +21,8 @@
 parser.add_argument('--no_normalize_features', action='store_true')
 parser.add_argument('--K', type=int, default=10)
 parser.add_argument('--alpha', type=float, default=0.1)
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 
@@ -47,4 +50,8 @@ def forward(self, data):
 dataset = get_planetoid_dataset(args.dataset, not args.no_normalize_features)
 permute_masks = random_planetoid_splits if args.random_splits else None
 run(dataset, Net(dataset), args.runs, args.epochs, args.lr, args.weight_decay,
-    args.early_stopping, permute_masks)
+    args.early_stopping, args.inference, args.profile, permute_masks)
+
+if args.profile:
+    rename_profile_file('citation', APPNP.__name__, args.dataset,
+                        str(args.random_splits))

benchmark/citation/arma.py

@@ -5,6 +5,7 @@
 from citation import get_planetoid_dataset, random_planetoid_splits, run
 
 from torch_geometric.nn import ARMAConv
+from torch_geometric.profile import rename_profile_file
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, required=True)
@@ -21,6 +22,8 @@
 parser.add_argument('--num_layers', type=int, default=1)
 parser.add_argument('--shared_weights', action='store_true')
 parser.add_argument('--skip_dropout', type=float, default=0.75)
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 
@@ -49,4 +52,8 @@ def forward(self, data):
 dataset = get_planetoid_dataset(args.dataset, not args.no_normalize_features)
 permute_masks = random_planetoid_splits if args.random_splits else None
 run(dataset, Net(dataset), args.runs, args.epochs, args.lr, args.weight_decay,
-    args.early_stopping, permute_masks)
+    args.early_stopping, args.inference, args.profile, permute_masks)
+
+if args.profile:
+    rename_profile_file('citation', ARMAConv.__name__, args.dataset,
+                        str(args.random_splits))

benchmark/citation/cheb.py

@@ -5,6 +5,7 @@
 from citation import get_planetoid_dataset, random_planetoid_splits, run
 
 from torch_geometric.nn import ChebConv
+from torch_geometric.profile import rename_profile_file
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, required=True)
@@ -18,6 +19,8 @@
 parser.add_argument('--dropout', type=float, default=0.5)
 parser.add_argument('--no_normalize_features', action='store_true')
 parser.add_argument('--num_hops', type=int, default=3)
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 
@@ -42,4 +45,8 @@ def forward(self, data):
 dataset = get_planetoid_dataset(args.dataset, not args.no_normalize_features)
 permute_masks = random_planetoid_splits if args.random_splits else None
 run(dataset, Net(dataset), args.runs, args.epochs, args.lr, args.weight_decay,
-    args.early_stopping, permute_masks)
+    args.early_stopping, args.inference, args.profile, permute_masks)
+
+if args.profile:
+    rename_profile_file('citation', ChebConv.__name__, args.dataset,
+                        str(args.random_splits))

benchmark/citation/gat.py

@@ -5,6 +5,7 @@
 from citation import get_planetoid_dataset, random_planetoid_splits, run
 
 from torch_geometric.nn import GATConv
+from torch_geometric.profile import rename_profile_file
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, required=True)
@@ -19,6 +20,8 @@
 parser.add_argument('--no_normalize_features', action='store_true')
 parser.add_argument('--heads', type=int, default=8)
 parser.add_argument('--output_heads', type=int, default=1)
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 
@@ -47,4 +50,8 @@ def forward(self, data):
 dataset = get_planetoid_dataset(args.dataset, not args.no_normalize_features)
 permute_masks = random_planetoid_splits if args.random_splits else None
 run(dataset, Net(dataset), args.runs, args.epochs, args.lr, args.weight_decay,
-    args.early_stopping, permute_masks)
+    args.early_stopping, args.inference, args.profile, permute_masks)
+
+if args.profile:
+    rename_profile_file('citation', GATConv.__name__, args.dataset,
+                        str(args.random_splits))

benchmark/citation/gcn.py

@@ -5,6 +5,7 @@
 from citation import get_planetoid_dataset, random_planetoid_splits, run
 
 from torch_geometric.nn import GCNConv
+from torch_geometric.profile import rename_profile_file
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, required=True)
@@ -17,6 +18,8 @@
 parser.add_argument('--hidden', type=int, default=16)
 parser.add_argument('--dropout', type=float, default=0.5)
 parser.add_argument('--no_normalize_features', action='store_true')
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 
@@ -41,4 +44,8 @@ def forward(self, data):
 dataset = get_planetoid_dataset(args.dataset, not args.no_normalize_features)
 permute_masks = random_planetoid_splits if args.random_splits else None
 run(dataset, Net(dataset), args.runs, args.epochs, args.lr, args.weight_decay,
-    args.early_stopping, permute_masks)
+    args.early_stopping, args.inference, args.profile, permute_masks)
+
+if args.profile:
+    rename_profile_file('citation', GCNConv.__name__, args.dataset,
+                        str(args.random_splits))

benchmark/citation/inference.sh

@@ -0,0 +1,118 @@
+#!/bin/sh
+
+echo "Cora"
+echo "===="
+
+echo "GCN"
+python gcn.py --dataset=Cora --inference
+python gcn.py --dataset=Cora --random_splits --inference
+python gcn.py --dataset=Cora --inference --profile
+python gcn.py --dataset=Cora --random_splits --inference --profile
+
+echo "GAT"
+python gat.py --dataset=Cora --inference
+python gat.py --dataset=Cora --random_splits --inference
+python gat.py --dataset=Cora --inference --profile
+python gat.py --dataset=Cora --random_splits --inference --profile
+
+echo "Cheby"
+python cheb.py --dataset=Cora --num_hops=3 --inference
+python cheb.py --dataset=Cora --num_hops=3 --random_splits --inference
+python cheb.py --dataset=Cora --num_hops=3 --inference --profile
+python cheb.py --dataset=Cora --num_hops=3 --random_splits --inference --profile
+
+echo "SGC"
+python sgc.py --dataset=Cora --K=3 --weight_decay=0.0005 --inference
+python sgc.py --dataset=Cora --K=3 --weight_decay=0.0005 --random_splits --inference
+python sgc.py --dataset=Cora --K=3 --weight_decay=0.0005 --inference --profile
+python sgc.py --dataset=Cora --K=3 --weight_decay=0.0005 --random_splits --inference --profile
+
+echo "ARMA"
+python arma.py --dataset=Cora --num_stacks=2 --num_layers=1 --shared_weights=True --inference
+python arma.py --dataset=Cora --num_stacks=3 --num_layers=1 --shared_weights=True --random_splits --inference
+python arma.py --dataset=Cora --num_stacks=2 --num_layers=1 --shared_weights=True --inference --profile
+python arma.py --dataset=Cora --num_stacks=3 --num_layers=1 --shared_weights=True --random_splits --inference --profile
+
+echo "APPNP"
+python appnp.py --dataset=Cora --alpha=0.1 --inference
+python appnp.py --dataset=Cora --alpha=0.1 --random_splits --inference
+python appnp.py --dataset=Cora --alpha=0.1 --inference --profile
+python appnp.py --dataset=Cora --alpha=0.1 --random_splits --inference --profile
+
+echo "CiteSeer"
+echo "========"
+
+echo "GCN"
+python gcn.py --dataset=CiteSeer --inference
+python gcn.py --dataset=CiteSeer --random_splits --inference
+python gcn.py --dataset=CiteSeer --inference --profile
+python gcn.py --dataset=CiteSeer --random_splits --inference --profile
+
+echo "GAT"
+python gat.py --dataset=CiteSeer --inference
+python gat.py --dataset=CiteSeer --random_splits --inference
+python gat.py --dataset=CiteSeer --inference --profile
+python gat.py --dataset=CiteSeer --random_splits --inference --profile
+
+echo "Cheby"
+python cheb.py --dataset=CiteSeer --num_hops=2 --inference
+python cheb.py --dataset=CiteSeer --num_hops=3 --random_splits --inference
+python cheb.py --dataset=CiteSeer --num_hops=2 --inference --profile
+python cheb.py --dataset=CiteSeer --num_hops=3 --random_splits --inference --profile
+
+echo "SGC"
+python sgc.py --dataset=CiteSeer --K=2 --weight_decay=0.005 --inference
+python sgc.py --dataset=CiteSeer --K=2 --weight_decay=0.005 --random_splits --inference
+python sgc.py --dataset=CiteSeer --K=2 --weight_decay=0.005 --inference --profile
+python sgc.py --dataset=CiteSeer --K=2 --weight_decay=0.005 --random_splits --inference --profile
+
+echo "ARMA"
+python arma.py --dataset=CiteSeer --num_stacks=3 --num_layers=1 --shared_weights=True --inference
+python arma.py --dataset=CiteSeer --num_stacks=3 --num_layers=1 --shared_weights=True --random_splits --inference
+python arma.py --dataset=CiteSeer --num_stacks=3 --num_layers=1 --shared_weights=True --inference --profile
+python arma.py --dataset=CiteSeer --num_stacks=3 --num_layers=1 --shared_weights=True --random_splits --inference --profile
+
+echo "APPNP"
+python appnp.py --dataset=CiteSeer --alpha=0.1 --inference
+python appnp.py --dataset=CiteSeer --alpha=0.1 --random_splits --inference
+python appnp.py --dataset=CiteSeer --alpha=0.1 --inference --profile
+python appnp.py --dataset=CiteSeer --alpha=0.1 --random_splits --inference --profile
+
+echo "PubMed"
+echo "======"
+
+echo "GCN"
+python gcn.py --dataset=PubMed --inference
+python gcn.py --dataset=PubMed --random_splits --inference
+python gcn.py --dataset=PubMed --inference --profile
+python gcn.py --dataset=PubMed --random_splits --inference --profile
+
+echo "GAT"
+python gat.py --dataset=PubMed --lr=0.01 --weight_decay=0.001 --output_heads=8 --inference
+python gat.py --dataset=PubMed --lr=0.01 --weight_decay=0.001 --output_heads=8 --random_splits --inference
+python gat.py --dataset=PubMed --lr=0.01 --weight_decay=0.001 --output_heads=8 --inference --profile
+python gat.py --dataset=PubMed --lr=0.01 --weight_decay=0.001 --output_heads=8 --random_splits --inference --profile
+
+echo "Cheby"
+python cheb.py --dataset=PubMed --num_hops=2 --inference
+python cheb.py --dataset=PubMed --num_hops=2 --random_splits --inference
+python cheb.py --dataset=PubMed --num_hops=2 --inference --profile
+python cheb.py --dataset=PubMed --num_hops=2 --random_splits --inference --profile
+
+echo "SGC"
+python sgc.py --dataset=PubMed --K=2 --weight_decay=0.0005 --inference
+python sgc.py --dataset=PubMed --K=2 --weight_decay=0.0005 --random_splits --inference
+python sgc.py --dataset=PubMed --K=2 --weight_decay=0.0005 --inference --profile
+python sgc.py --dataset=PubMed --K=2 --weight_decay=0.0005 --random_splits --inference --profile
+
+echo "ARMA"
+python arma.py --dataset=PubMed --num_stacks=2 --num_layers=1 --skip_dropout=0 --inference
+python arma.py --dataset=PubMed --num_stacks=2 --num_layers=1 --skip_dropout=0.5 --random_splits --inference
+python arma.py --dataset=PubMed --num_stacks=2 --num_layers=1 --skip_dropout=0 --inference --profile
+python arma.py --dataset=PubMed --num_stacks=2 --num_layers=1 --skip_dropout=0.5 --random_splits --inference --profile
+
+echo "APPNP"
+python appnp.py --dataset=PubMed --alpha=0.1 --inference
+python appnp.py --dataset=PubMed --alpha=0.1 --random_splits --inference
+python appnp.py --dataset=PubMed --alpha=0.1 --inference --profile
+python appnp.py --dataset=PubMed --alpha=0.1 --random_splits --inference --profile

benchmark/citation/sgc.py

@@ -5,6 +5,7 @@
 from citation import get_planetoid_dataset, random_planetoid_splits, run
 
 from torch_geometric.nn import SGConv
+from torch_geometric.profile import rename_profile_file
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dataset', type=str, required=True)
@@ -16,6 +17,8 @@
 parser.add_argument('--early_stopping', type=int, default=10)
 parser.add_argument('--no_normalize_features', action='store_true')
 parser.add_argument('--K', type=int, default=2)
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 
@@ -37,4 +40,8 @@ def forward(self, data):
 dataset = get_planetoid_dataset(args.dataset, not args.no_normalize_features)
 permute_masks = random_planetoid_splits if args.random_splits else None
 run(dataset, Net(dataset), args.runs, args.epochs, args.lr, args.weight_decay,
-    args.early_stopping, permute_masks)
+    args.early_stopping, args.inference, args.profile, permute_masks)
+
+if args.profile:
+    rename_profile_file('citation', SGConv.__name__, args.dataset,
+                        str(args.random_splits))

benchmark/citation/train_eval.py

@@ -4,7 +4,9 @@
 import torch.nn.functional as F
 from torch import tensor
 from torch.optim import Adam
+from torch.profiler import ProfilerActivity, profile
 
+from torch_geometric.profile import trace_handler
 from torch_geometric.utils import index_to_mask
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -34,9 +36,8 @@ def random_planetoid_splits(data, num_classes):
     return data
 
 
-def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
-        permute_masks=None, logger=None):
-
+def run_train(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
+              permute_masks=None, logger=None):
     val_losses, accs, durations = [], [], []
     for _ in range(runs):
         data = dataset[0]
@@ -82,12 +83,52 @@ def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
         val_losses.append(best_val_loss)
         accs.append(test_acc)
         durations.append(t_end - t_start)
-
     loss, acc, duration = tensor(val_losses), tensor(accs), tensor(durations)
 
     print(f'Val Loss: {float(loss.mean()):.4f}, '
           f'Test Accuracy: {float(acc.mean()):.3f} ± {float(acc.std()):.3f}, '
-          f'Duration: {float(duration.mean()):.3f}')
+          f'Duration: {float(duration.mean()):.3f}s')
+
+
+@torch.no_grad()
+def run_inference(dataset, model, epochs, profiling, permute_masks=None,
+                  logger=None):
+    data = dataset[0]
+    if permute_masks is not None:
+        data = permute_masks(data, dataset.num_classes)
+    data = data.to(device)
+
+    model.to(device).reset_parameters()
+
+    for epoch in range(1, epochs + 1):
+        if epoch == epochs:
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
+            t_start = time.time()
+
+        inference(model, data)
+
+        if epoch == epochs:
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
+            t_end = time.time()
+            duration = t_end - t_start
+            print(f'End-to-End Inference Time: {duration:.8f}s', flush=True)
+
+    if profiling:
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+                     on_trace_ready=trace_handler) as p:
+            inference(model, data)
+            p.step()
+
+
+def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
+        inference, profiling, permute_masks=None, logger=None):
+    if not inference:
+        run_train(dataset, model, runs, epochs, lr, weight_decay,
+                  early_stopping, permute_masks, logger)
+    else:
+        run_inference(dataset, model, epochs, profiling, permute_masks, logger)
 
 
 def train(model, optimizer, data):
@@ -116,3 +157,9 @@ def evaluate(model, data):
         outs[f'{key}_acc'] = acc
 
     return outs
+
+
+@torch.no_grad()
+def inference(model, data):
+    model.eval()
+    model(data)