Update and add log argument in timeit

pyg-team · Aug 24, 2022 · 696a89a · 696a89a
1 parent 4fef86a
commit 696a89a
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 13 deletions.
diff --git a/benchmark/inference/inference_benchmark.py b/benchmark/inference/inference_benchmark.py
@@ -10,8 +10,8 @@
 
 supported_sets = {
     'ogbn-mag': ['rgat', 'rgcn'],
-    'ogbn-products': ['edge_cnn', 'gat', 'gcn', 'pna', 'graph_sage'],
-    'Reddit': ['edge_cnn', 'gat', 'gcn', 'pna', 'graph_sage'],
+    'ogbn-products': ['edge_cnn', 'gat', 'gcn', 'pna', 'sage'],
+    'Reddit': ['edge_cnn', 'gat', 'gcn', 'pna', 'sage'],
 }
 
 

diff --git a/benchmark/inference/utils.py b/benchmark/inference/utils.py
@@ -19,7 +19,7 @@
     'gat': GAT,
     'gcn': GCN,
     'pna': PNA,
-    'graph_sage': GraphSAGE,
+    'sage': GraphSAGE,
     'rgat': HeteroGAT,
     'rgcn': HeteroGraphSAGE,
 }

diff --git a/benchmark/kernel/main_performance.py b/benchmark/kernel/main_performance.py
@@ -70,8 +70,8 @@ def run_train():
             dataset_name)
 
         for num_layers, hidden in product(layers, hiddens):
-            print("--\n{} - {} - {} - {}".format(dataset_name, Net.__name__,
-                                                 num_layers, hidden))
+            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)
@@ -80,10 +80,10 @@ def run_train():
             acc_list = []
             for epoch in range(1, args.epochs + 1):
                 loss, stats = train(model, optimizer, train_loader)
-                with timeit() as t:
+                with timeit(log=False) as t:
                     val_acc = eval_acc(model, val_loader)
                 val_time = t.duration
-                with timeit() as t:
+                with timeit(log=False) as t:
                     test_acc = eval_acc(model, test_loader)
                 test_time = t.duration
 
@@ -101,8 +101,8 @@ def run_inference():
         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))
+            print(f'--\n{dataset_name} - {Net.__name__}'
+                  '- {num_layers} - {hidden}')
 
             model = Net(dataset, num_layers, hidden).to(device)
 

diff --git a/torch_geometric/profile/profile.py b/torch_geometric/profile/profile.py
@@ -96,7 +96,7 @@ def wrapper(*args, **kwargs) -> Tuple[Any, Stats]:
 
 
 class timeit(ContextDecorator):
-    r"""A ContextDecorator to facilitate timing a function, *e.g.*, obtaining
+    r"""A context decorator to facilitate timing a function, *e.g.*, obtaining
     the runtime of a specific model on a specific dataset.
 
     .. code-block:: python
@@ -105,10 +105,13 @@ class timeit(ContextDecorator):
         def test(model, x, edge_index):
             return model(x, edge_index)
 
-        with timeit() as timeit:
+        with timeit() as t:
             z = test(model, x, edge_index)
-        time = timeit.duration
+        time = t.duration
     """
+    def __init__(self, log: bool = True):
+        self.log = log
+
     def __enter__(self):
         if torch.cuda.is_available():
             torch.cuda.synchronize()
@@ -120,7 +123,8 @@ def __exit__(self, *args):
             torch.cuda.synchronize()
         self.t_end = time.time()
         self.duration = self.t_end - self.t_start
-        print(f'Time: {self.duration:.8f}s', flush=True)
+        if self.log:
+            print(f'Time: {self.duration:.8f}s', flush=True)
 
 
 def get_stats_summary(stats_list: List[Stats]):