pyg-team · rusty1s · Sep 1, 2022 · Aug 31, 2022 · Sep 1, 2022 · Sep 1, 2022
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,7 +9,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))
 - Added `pyg::sampler::Mapper` utility for mapping global to local node indices ([#45](https://github.com/pyg-team/pyg-lib/pull/45)))
-- Added benchmark script ([#45](https://github.com/pyg-team/pyg-lib/pull/45))
+- Added benchmark script ([#45](https://github.com/pyg-team/pyg-lib/pull/45), [#79](https://github.com/pyg-team/pyg-lib/pull/79))
 - 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))

diff --git a/benchmark/sampler/neighbor.py b/benchmark/sampler/neighbor.py
@@ -0,0 +1,70 @@
+import argparse
+import time
+
+import torch
+
+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=int)
+argparser.add_argument('--replace', action='store_true')
+argparser.add_argument('--directed', action='store_true')
+
+args = argparser.parse_args()
+
+
+@withSeed
+@withDataset('DIMACS10', 'citationCiteseer')
+def test_neighbor(dataset, **kwargs):
+    (rowptr, col), num_nodes = dataset, dataset[0].size(0) - 1
+
+    for num_neighbors in args.num_neighbors:
+        for batch_size in args.batch_sizes:
+            # pyg-lib neighbor sampler
+            start = time.perf_counter()
+            for node in tqdm(range(0, num_nodes, batch_size)):
+                last_seed_node = node + batch_size \
+                            if node + batch_size < num_nodes else num_nodes
+                seed = torch.arange(node, last_seed_node)
+                pyg_lib.sampler.neighbor_sample(rowptr, col, seed,
+                                                num_neighbors, args.replace,
+                                                args.directed, disjoint=False,
+                                                return_edge_id=True)
+            stop = time.perf_counter()
+            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 node in tqdm(range(0, num_nodes, batch_size)):
+                last_seed_node = node + batch_size \
+                            if node + batch_size < num_nodes else num_nodes
+                seed = torch.arange(node, last_seed_node)
+                torch.ops.torch_sparse.neighbor_sample(rowptr, col, seed,
+                                                       num_neighbors,
+                                                       args.replace,
+                                                       args.directed)
+            stop = time.perf_counter()
+            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__':
+    test_neighbor()