Add profiling support for benchmark/kernel and benchmark/inference

benchmark/citation/train_eval.py

@@ -4,9 +4,8 @@
 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.profile import timeit, torch_profile
 from torch_geometric.utils import index_to_mask
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -102,24 +101,14 @@ def run_inference(dataset, model, epochs, profiling, permute_masks=None,
 
     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)
+            with timeit():
+                inference(model, data)
+        else:
+            inference(model, data)
 
     if profiling:
-        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-                     on_trace_ready=trace_handler) as p:
+        with torch_profile():
             inference(model, data)
-            p.step()
 
 
 def run(dataset, model, runs, epochs, lr, weight_decay, early_stopping,

benchmark/inference/inference_benchmark.py

@@ -1,16 +1,16 @@
 import argparse
-from timeit import default_timer
 
 import torch
 from utils import get_dataset, get_model
 
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import PNAConv
+from torch_geometric.profile import rename_profile_file, timeit, torch_profile
 
 supported_sets = {
     'ogbn-mag': ['rgat', 'rgcn'],
-    'ogbn-products': ['edge_cnn', 'gat', 'gcn', 'pna'],
-    'Reddit': ['edge_cnn', 'gat', 'gcn', 'pna'],
+    'ogbn-products': ['edge_cnn', 'gat', 'gcn', 'pna', 'graph_sage'],
+    'Reddit': ['edge_cnn', 'gat', 'gcn', 'pna', 'graph_sage'],
 }
 
 
@@ -92,11 +92,21 @@ def run(args: argparse.ArgumentParser) -> None:
                         model = model.to(device)
                         model.eval()
 
-                        start = default_timer()
-                        model.inference(subgraph_loader, device,
-                                        progress_bar=True)
-                        stop = default_timer()
-                        print(f'Inference time={stop-start:.3f} seconds\n')
+                        for _ in range(args.warmup):
+                            model.inference(subgraph_loader, device,
+                                            progress_bar=True)
+                        with timeit():
+                            model.inference(subgraph_loader, device,
+                                            progress_bar=True)
+
+                        if args.profile:
+                            with torch_profile():
+                                model.inference(subgraph_loader, device,
+                                                progress_bar=True)
+                            rename_profile_file(
+                                model_name, dataset_name, str(batch_size),
+                                str(layers), str(hidden_channels),
+                                str(subgraph_loader.num_neighbors))
 
 
 if __name__ == '__main__':
@@ -120,7 +130,9 @@ def run(args: argparse.ArgumentParser) -> None:
     argparser.add_argument(
         '--hetero-num-neighbors', default=-1, type=int,
         help='number of neighbors to sample per layer for hetero workloads')
-    argparser.add_argument('--num-workers', default=2, type=int)
+    argparser.add_argument('--num-workers', default=0, type=int)
+    argparser.add_argument('--warmup', default=1, type=int)
+    argparser.add_argument('--profile', action='store_true')
 
     args = argparser.parse_args()
 

benchmark/inference/utils.py

@@ -6,13 +6,20 @@
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import OGB_MAG, Reddit
-from torch_geometric.nn.models.basic_gnn import GAT, GCN, PNA, EdgeCNN
+from torch_geometric.nn.models.basic_gnn import (
+    GAT,
+    GCN,
+    PNA,
+    EdgeCNN,
+    GraphSAGE,
+)
 
 models_dict = {
     'edge_cnn': EdgeCNN,
     'gat': GAT,
     'gcn': GCN,
     'pna': PNA,
+    'graph_sage': GraphSAGE,
     'rgat': HeteroGAT,
     'rgcn': HeteroGraphSAGE,
 }

benchmark/kernel/main_performance.py

@@ -6,11 +6,17 @@
 from gcn import GCN
 from gin import GIN
 from graph_sage import GraphSAGE
-from train_eval import eval_acc, train
+from train_eval import eval_acc, inference_run, train
 
 from torch_geometric import seed_everything
 from torch_geometric.loader import DataLoader
-from torch_geometric.profile import get_stats_summary, profileit, timeit
+from torch_geometric.profile import (
+    get_stats_summary,
+    profileit,
+    rename_profile_file,
+    timeit,
+    torch_profile,
+)
 
 seed_everything(0)
 
@@ -22,6 +28,8 @@
                     help='Skip the first few runs')
 parser.add_argument('--goal_accuracy', type=int, default=1,
                     help='The goal test accuracy')
+parser.add_argument('--inference', action='store_true')
+parser.add_argument('--profile', action='store_true')
 args = parser.parse_args()
 
 layers = [1, 2, 3]
@@ -37,11 +45,8 @@
 
 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):
+def prepare_dataloader(dataset_name):
     dataset = get_dataset(dataset_name, sparse=True)
     num_train = int(len(dataset) * 0.8)
     num_val = int(len(dataset) * 0.1)
@@ -56,21 +61,68 @@
                             shuffle=False)
     test_loader = DataLoader(test_dataset, batch_size=args.batch_size,
                              shuffle=False)
-
-    for num_layers, hidden in product(layers, hiddens):
-        print(f'--\n{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)
+    return dataset, train_loader, val_loader, test_loader
+
+
+def run_train():
+    for dataset_name, Net in product(datasets, nets):
+        dataset, train_loader, val_loader, test_loader = prepare_dataloader(
+            dataset_name)
+
+        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 = []
+            acc_list = []
+            for epoch in range(1, args.epochs + 1):
+                loss, stats = train(model, optimizer, train_loader)
+                with timeit() as t:
+                    val_acc = eval_acc(model, val_loader)
+                val_time = t.duration
+                with timeit() as t:
+                    test_acc = eval_acc(model, test_loader)
+                test_time = t.duration
+
+                if epoch >= args.warmup_profile:
+                    stats_list.append(stats)
+                    acc_list.append([val_acc, val_time, test_acc, test_time])
+
+            stats_summary = get_stats_summary(stats_list)
+            print(stats_summary)
+
+
+@torch.no_grad()
+def run_inference():
+    for dataset_name, Net in product(datasets, nets):
+        dataset, _, _, test_loader = prepare_dataloader(dataset_name)
+
+        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)
+
+            for epoch in range(1, args.epochs + 1):
+                if epoch == args.epochs:
+                    with timeit():
+                        inference_run(model, test_loader)
+                else:
+                    inference_run(model, test_loader)
+
+            if args.profile:
+                with torch_profile():
+                    inference_run(model, test_loader)
+                rename_profile_file(Net.__name__, dataset_name,
+                                    str(num_layers), str(hidden))
+
+
+if not args.inference:
+    # Decorate train functions:
+    train = profileit()(train)
+    run_train()
+else:
+    run_inference()

benchmark/kernel/train_eval.py

@@ -143,3 +143,11 @@ def eval_loss(model, loader):
             out = model(data)
         loss += F.nll_loss(out, data.y.view(-1), reduction='sum').item()
     return loss / len(loader.dataset)
+
+
+def inference_run(model, loader):
+    model.eval()
+    for data in loader:
+        data = data.to(device)
+        with torch.no_grad():
+            model(data)

benchmark/points/train_eval.py

@@ -3,10 +3,9 @@
 import torch
 import torch.nn.functional as F
 from torch.optim import Adam
-from torch.profiler import ProfilerActivity, profile
 
 from torch_geometric.loader import DataLoader
-from torch_geometric.profile import trace_handler
+from torch_geometric.profile import timeit, torch_profile
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -47,25 +46,16 @@ def run_inference(test_dataset, model, epochs, batch_size, profiling):
     test_loader = DataLoader(test_dataset, batch_size, shuffle=False)
 
     for epoch in range(1, epochs + 1):
+        print("Epoch: ", epoch)
         if epoch == epochs:
-            if torch.cuda.is_available():
-                torch.cuda.synchronize()
-            t_start = time.time()
-
-        inference(model, test_loader, device)
-
-        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)
+            with timeit():
+                inference(model, test_loader, device)
+        else:
+            inference(model, test_loader, device)
 
     if profiling:
-        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-                     on_trace_ready=trace_handler) as p:
+        with torch_profile():
             inference(model, test_loader, device)
-            p.step()
 
 
 def run(train_dataset, test_dataset, model, epochs, batch_size, lr,

test/profile/test_profile.py

@@ -2,16 +2,15 @@
 
 import torch
 import torch.nn.functional as F
-from torch.profiler import ProfilerActivity, profile
 
 from torch_geometric.nn import GraphSAGE
 from torch_geometric.profile import (
     get_stats_summary,
     profileit,
     rename_profile_file,
     timeit,
-    trace_handler,
 )
+from torch_geometric.profile.profile import torch_profile
 from torch_geometric.testing import onlyFullTest, withCUDA
 
 
@@ -69,18 +68,24 @@ def test(model, x, edge_index, y):
 
 
 @onlyFullTest
-def test_trace_handler(get_dataset):
+def test_torch_profile(get_dataset):
     dataset = get_dataset(name='PubMed')
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     data = dataset[0].to(device)
     model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
                       out_channels=dataset.num_classes).to(device)
     model.eval()
 
+    @timeit()
+    def inference_e2e(model, data):
+        model(data.x, data.edge_index)
+
+    @torch_profile()
+    def inference_profile(model, data):
+        model(data.x, data.edge_index)
+
     for epoch in range(3):
-        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-                     on_trace_ready=trace_handler) as p:
-            model(data.x, data.edge_index)
-            p.step()
+        inference_e2e(model, data)
+        inference_profile(model, data)
     rename_profile_file('test_profile')
     assert os.path.exists('profile-test_profile.json')

torch_geometric/profile/__init__.py

@@ -1,5 +1,5 @@
 from .profile import profileit, timeit, get_stats_summary
-from .profile import trace_handler, rename_profile_file
+from .profile import trace_handler, rename_profile_file, torch_profile
 from .utils import count_parameters
 from .utils import get_model_size
 from .utils import get_data_size
@@ -13,6 +13,7 @@
     'get_stats_summary',
     'trace_handler',
     'rename_profile_file',
+    'torch_profile',
     'count_parameters',
     'get_model_size',
     'get_data_size',