Merge branch 'yanbing/benchmark_inference' of https://github.com/yanb…

…ing-j/pytorch_geometric into yanbing/benchmark_inference

benchmark/citation/appnp.py

@@ -21,7 +21,8 @@
 parser.add_argument('--K', type=int, default=10)
 parser.add_argument('--alpha', type=float, default=0.1)
 parser.add_argument('--inference', type=bool, default=False)
-parser.add_argument('--profile', type=bool, default=False) # Currently support profile in inference
+parser.add_argument('--profile', type=bool,
+                    default=False)  # Currently support profile in inference
 args = parser.parse_args()
 
 
@@ -56,5 +57,6 @@ def forward(self, data):
     import os
     import pathlib
     profile_dir = str(pathlib.Path.cwd()) + '/'
-    timeline_file = profile_dir + 'profile-citation-APPNP-' + args.dataset + '-random_splits-' + str(args.random_splits) + '.json'
+    timeline_file = profile_dir + 'profile-citation-APPNP-' + args.dataset + '-random_splits-' + str(
+        args.random_splits) + '.json'
     os.rename('timeline.json', timeline_file)

benchmark/citation/arma.py

@@ -22,7 +22,8 @@
 parser.add_argument('--shared_weights', type=bool, default=False)
 parser.add_argument('--skip_dropout', type=float, default=0.75)
 parser.add_argument('--inference', type=bool, default=False)
-parser.add_argument('--profile', type=bool, default=False) # Currently support profile in inference
+parser.add_argument('--profile', type=bool,
+                    default=False)  # Currently support profile in inference
 args = parser.parse_args()
 
 
@@ -58,5 +59,6 @@ def forward(self, data):
     import os
     import pathlib
     profile_dir = str(pathlib.Path.cwd()) + '/'
-    timeline_file = profile_dir + 'profile-citation-ARMA-' + args.dataset + '-random_splits-' + str(args.random_splits) + '.json'
+    timeline_file = profile_dir + 'profile-citation-ARMA-' + args.dataset + '-random_splits-' + str(
+        args.random_splits) + '.json'
     os.rename('timeline.json', timeline_file)

benchmark/citation/cheb.py

@@ -19,7 +19,8 @@
 parser.add_argument('--normalize_features', type=bool, default=True)
 parser.add_argument('--num_hops', type=int, default=3)
 parser.add_argument('--inference', type=bool, default=False)
-parser.add_argument('--profile', type=bool, default=False) # Currently support profile in inference
+parser.add_argument('--profile', type=bool,
+                    default=False)  # Currently support profile in inference
 args = parser.parse_args()
 
 
@@ -51,5 +52,6 @@ def forward(self, data):
     import os
     import pathlib
     profile_dir = str(pathlib.Path.cwd()) + '/'
-    timeline_file = profile_dir + 'profile-citation-CHEBY-' + args.dataset + '-random_splits-' + str(args.random_splits) + '.json'
+    timeline_file = profile_dir + 'profile-citation-CHEBY-' + args.dataset + '-random_splits-' + str(
+        args.random_splits) + '.json'
     os.rename('timeline.json', timeline_file)

benchmark/citation/gat.py

@@ -20,7 +20,8 @@
 parser.add_argument('--heads', type=int, default=8)
 parser.add_argument('--output_heads', type=int, default=1)
 parser.add_argument('--inference', type=bool, default=False)
-parser.add_argument('--profile', type=bool, default=False) # Currently support profile in inference
+parser.add_argument('--profile', type=bool,
+                    default=False)  # Currently support profile in inference
 args = parser.parse_args()
 
 
@@ -56,5 +57,6 @@ def forward(self, data):
     import os
     import pathlib
     profile_dir = str(pathlib.Path.cwd()) + '/'
-    timeline_file = profile_dir + 'profile-citation-GAT-' + args.dataset + '-random_splits-' + str(args.random_splits) + '.json'
+    timeline_file = profile_dir + 'profile-citation-GAT-' + args.dataset + '-random_splits-' + str(
+        args.random_splits) + '.json'
     os.rename('timeline.json', timeline_file)

benchmark/citation/gcn.py

@@ -18,7 +18,8 @@
 parser.add_argument('--dropout', type=float, default=0.5)
 parser.add_argument('--normalize_features', type=bool, default=True)
 parser.add_argument('--inference', type=bool, default=False)
-parser.add_argument('--profile', type=bool, default=False) # Currently support profile in inference
+parser.add_argument('--profile', type=bool,
+                    default=False)  # Currently support profile in inference
 args = parser.parse_args()
 
 
@@ -50,5 +51,6 @@ def forward(self, data):
     import os
     import pathlib
     profile_dir = str(pathlib.Path.cwd()) + '/'
-    timeline_file = profile_dir + 'profile-citation-GCN-' + args.dataset + '-random_splits-' + str(args.random_splits) + '.json'
+    timeline_file = profile_dir + 'profile-citation-GCN-' + args.dataset + '-random_splits-' + str(
+        args.random_splits) + '.json'
     os.rename('timeline.json', timeline_file)

benchmark/citation/sgc.py

@@ -17,7 +17,8 @@
 parser.add_argument('--normalize_features', type=bool, default=False)
 parser.add_argument('--K', type=int, default=2)
 parser.add_argument('--inference', type=bool, default=False)
-parser.add_argument('--profile', type=bool, default=False) # Currently support profile in inference
+parser.add_argument('--profile', type=bool,
+                    default=False)  # Currently support profile in inference
 args = parser.parse_args()
 
 
@@ -46,5 +47,6 @@ def forward(self, data):
     import os
     import pathlib
     profile_dir = str(pathlib.Path.cwd()) + '/'
-    timeline_file = profile_dir + 'profile-citation-SGC-' + args.dataset + '-random_splits-' + str(args.random_splits) + '.json'
+    timeline_file = profile_dir + 'profile-citation-SGC-' + args.dataset + '-random_splits-' + str(
+        args.random_splits) + '.json'
     os.rename('timeline.json', timeline_file)

benchmark/citation/train_eval.py

@@ -4,12 +4,13 @@
 import torch.nn.functional as F
 from torch import tensor
 from torch.optim import Adam
-from torch.profiler import profile, ProfilerActivity
+from torch.profiler import ProfilerActivity, profile
 
 from torch_geometric.utils import index_to_mask
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-profile_sort = "self_cuda_time_total" if torch.cuda.is_available() else "self_cpu_time_total"
+profile_sort = "self_cuda_time_total" if torch.cuda.is_available(
+) else "self_cpu_time_total"
 
 
 def random_planetoid_splits(data, num_classes):
@@ -35,6 +36,7 @@ def random_planetoid_splits(data, num_classes):
 
     return data
 
+
 def trace_handler(p):
     output = p.key_averages().table(sort_by=profile_sort)
     print(output)
@@ -43,8 +45,9 @@ def trace_handler(p):
     timeline_file = profile_dir + 'timeline' + '.json'
     p.export_chrome_trace(timeline_file)
 
-def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping, inference, profiling,
-        permute_masks=None, logger=None):
+
+def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping,
+        inference, profiling, permute_masks=None, logger=None):
     val_losses, accs, durations = [], [], []
     if not inference:
         for _ in range(runs):
@@ -54,7 +57,8 @@ def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping, inferenc
             data = data.to(device)
 
             model.to(device).reset_parameters()
-            optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+            optimizer = Adam(model.parameters(), lr=lr,
+                             weight_decay=weight_decay)
 
             if torch.cuda.is_available():
                 torch.cuda.synchronize()
@@ -91,9 +95,11 @@ def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping, inferenc
             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)
+        loss, acc, duration = tensor(val_losses), tensor(accs), tensor(
+            durations)
 
-        print(f'Val Loss: {float(loss.mean()):.4f}, '
+        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}s')
     else:
@@ -108,9 +114,10 @@ def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping, inferenc
             for epoch in range(1, epochs + 1):
                 if i == int(runs / 2) and epoch == int(epochs / 2):
                     if profiling:
-                        with profile(activities=[
-                            ProfilerActivity.CPU, ProfilerActivity.CUDA],
-                            on_trace_ready=trace_handler) as p:
+                        with profile(
+                                activities=[
+                                    ProfilerActivity.CPU, ProfilerActivity.CUDA
+                                ], on_trace_ready=trace_handler) as p:
                             test(model, data)
                             p.step()
                     else:
@@ -124,10 +131,12 @@ def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping, inferenc
                             torch.cuda.synchronize()
                         t_end = time.time()
                         duration = t_end - t_start
-                        print("End-to-End time: {} s".format(duration), flush=True)
+                        print("End-to-End time: {} s".format(duration),
+                              flush=True)
                 else:
                     test(model, data)
 
+
 def train(model, optimizer, data):
     model.train()
     optimizer.zero_grad()
@@ -155,6 +164,7 @@ def evaluate(model, data):
 
     return outs
 
+
 def test(model, data):
     model.eval()
     with torch.no_grad():