Update benchmark scripts

CHANGELOG.md

@@ -11,7 +11,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `grouped_matmul` and `segment_matmul` CUDA implementations via `cutlass` ([#51](https://github.com/pyg-team/pyg-lib/pull/51), [#56](https://github.com/pyg-team/pyg-lib/pull/56), [#61](https://github.com/pyg-team/pyg-lib/pull/61), [#64](https://github.com/pyg-team/pyg-lib/pull/64), [#69](https://github.com/pyg-team/pyg-lib/pull/69))
 - Added `pyg::sampler::neighbor_sample` implementation ([#54](https://github.com/pyg-team/pyg-lib/pull/54), [#76](https://github.com/pyg-team/pyg-lib/pull/76), [#77](https://github.com/pyg-team/pyg-lib/pull/77), [#78](https://github.com/pyg-team/pyg-lib/pull/78), [#80](https://github.com/pyg-team/pyg-lib/pull/80), [#81](https://github.com/pyg-team/pyg-lib/pull/81)), [#85](https://github.com/pyg-team/pyg-lib/pull/85), [#86](https://github.com/pyg-team/pyg-lib/pull/86), [#87](https://github.com/pyg-team/pyg-lib/pull/87), [#89](https://github.com/pyg-team/pyg-lib/pull/89))
 - Added `pyg::sampler::Mapper` utility for mapping global to local node indices ([#45](https://github.com/pyg-team/pyg-lib/pull/45), [#83](https://github.com/pyg-team/pyg-lib/pull/83))
-- Added benchmark script ([#45](https://github.com/pyg-team/pyg-lib/pull/45), [#79](https://github.com/pyg-team/pyg-lib/pull/79), [#82](https://github.com/pyg-team/pyg-lib/pull/82))
+- Added benchmark script ([#45](https://github.com/pyg-team/pyg-lib/pull/45), [#79](https://github.com/pyg-team/pyg-lib/pull/79), [#82](https://github.com/pyg-team/pyg-lib/pull/82), [#91](https://github.com/pyg-team/pyg-lib/pull/91))
 - Added download script for benchmark data ([#44](https://github.com/pyg-team/pyg-lib/pull/44))
 - Added `biased sampling` utils ([#38](https://github.com/pyg-team/pyg-lib/pull/38))
 - Added `CHANGELOG.md` ([#39](https://github.com/pyg-team/pyg-lib/pull/39))

benchmark/sampler/neighbor.py

@@ -3,27 +3,33 @@
 import time
 
 import torch
+from tqdm import tqdm
+
+import pyg_lib
+from pyg_lib.testing import withDataset, withSeed
 
 try:
     import torch_sparse  # noqa
     baseline_neighbor_sample = torch.ops.torch_sparse.neighbor_sample
 except ImportError:
     baseline_neighbor_sample = None
-from tqdm import tqdm
 
-import pyg_lib
-from pyg_lib.testing import withDataset, withSeed
-
-argparser = argparse.ArgumentParser('Neighbor Sampler benchmark')
-argparser.add_argument('--batch-sizes', nargs='+',
-                       default=[512, 1024, 2048, 4096, 8192], type=int)
-argparser.add_argument('--num_neighbors', default=[[-1], [15, 10, 5],
-                                                   [20, 15, 10]],
-                       type=ast.literal_eval)
+argparser = argparse.ArgumentParser()
+argparser.add_argument('--batch-sizes', nargs='+', type=int, default=[
+    512,
+    1024,
+    2048,
+    4096,
+    8192,
+])
+argparser.add_argument('--num_neighbors', type=ast.literal_eval, default=[
+    [-1],
+    [15, 10, 5],
+    [20, 15, 10],
+])
 argparser.add_argument('--replace', action='store_true')
 argparser.add_argument('--directed', action='store_true')
 argparser.add_argument('--shuffle', action='store_true')
-
 args = argparser.parse_args()
 
 
@@ -32,39 +38,44 @@
 def test_neighbor(dataset, **kwargs):
     (rowptr, col), num_nodes = dataset, dataset[0].size(0) - 1
 
+    if args.shuffle:
+        node_perm = torch.randperm(num_nodes)
+    else:
+        node_perm = torch.arange(num_nodes)
+
     for num_neighbors in args.num_neighbors:
         for batch_size in args.batch_sizes:
-            # pyg-lib neighbor sampler
-            start = time.perf_counter()
-            nodes_ids = torch.randperm(
-                num_nodes) if args.shuffle else torch.arange(0, num_nodes)
-            for seed in tqdm(nodes_ids.split(batch_size)):
-                pyg_lib.sampler.neighbor_sample(rowptr, col, seed,
-                                                num_neighbors,
-                                                replace=args.replace,
-                                                directed=args.directed,
-                                                disjoint=False,
-                                                return_edge_id=True)
-            stop = time.perf_counter()
+            print(f'pyg-lib      (batch_size={batch_size}, '
+                  f'num_neighbors={num_neighbors}):')
+            t = time.perf_counter()
+            for seed in tqdm(node_perm.split(batch_size)):
+                pyg_lib.sampler.neighbor_sample(
+                    rowptr,
+                    col,
+                    seed,
+                    num_neighbors,
+                    replace=args.replace,
+                    directed=args.directed,
+                    disjoint=False,
+                    return_edge_id=True,
+                )
+            print(f'time={time.perf_counter()-t:.3f} seconds')
             print('-------------------------')
-            print('pyg-lib neighbor sample')
-            print(f'Batch size={batch_size}, '
-                  f'Num_neighbors={num_neighbors}, '
-                  f'Time={stop-start:.3f} seconds\n')
 
-            # pytorch-sparse neighbor sampler
-            start = time.perf_counter()
-            for seed in tqdm(nodes_ids.split(batch_size)):
-                torch.ops.torch_sparse.neighbor_sample(rowptr, col, seed,
-                                                       num_neighbors,
-                                                       args.replace,
-                                                       args.directed)
-            stop = time.perf_counter()
+            print(f'torch-sparse (batch_size={batch_size}, '
+                  f'num_neighbors={num_neighbors}):')
+            t = time.perf_counter()
+            for seed in tqdm(node_perm.split(batch_size)):
+                torch.ops.torch_sparse.neighbor_sample(
+                    rowptr,
+                    col,
+                    seed,
+                    num_neighbors,
+                    args.replace,
+                    args.directed,
+                )
+            print(f'time={time.perf_counter()-t:.3f} seconds')
             print('-------------------------')
-            print('pytorch_sparse neighbor sample')
-            print(f'Batch size={batch_size}, '
-                  f'Num_neighbors={num_neighbors}, '
-                  f'Time={stop-start:.3f} seconds\n')
 
 
 if __name__ == '__main__':

benchmark/main.py → benchmark/sampler/subgraph.py

@@ -16,15 +16,15 @@ def test_subgraph(dataset, **kwargs):
     t = time.perf_counter()
     for _ in range(10):
         pyg_lib.sampler.subgraph(rowptr, col, nodes)
-    print(time.perf_counter() - t)
+    print(f'time={time.perf_counter()-t:.6f} seconds')
 
     edge_index = to_edge_index(rowptr, col)
     from torch_geometric.utils import subgraph
 
     t = time.perf_counter()
     for _ in range(10):
         subgraph(nodes, edge_index, num_nodes=num_nodes, relabel_nodes=True)
-    print(time.perf_counter() - t)
+    print(f'time={time.perf_counter()-t:.6f} seconds')
 
 
 if __name__ == '__main__':

pyg_lib/sampler/__init__.py

@@ -72,7 +72,7 @@ def hetero_neighbor_sample(
 
     .. note ::
         Similar to :meth:`neighbor_sample`, but expects a dictionary of node
-        types (:obj:`str`) and  edge tpyes (:obj:`Tuple[str, str, str]`) for
+        types (:obj:`str`) and  edge types (:obj:`Tuple[str, str, str]`) for
         each non-boolean argument.
 
     Args: