Skip to content

Commit

Permalink
[Benchmark] Benchmark script for KNN Graph (dmlc#3080)
Browse files Browse the repository at this point in the history
* add knn benchmark result

* update

* update knn benchmark to test on multiple algorithms

* update benchmark script

* update
  • Loading branch information
hetong007 authored Jul 1, 2021
1 parent a97aaa4 commit 1bfc311
Showing 1 changed file with 11 additions and 10 deletions.
21 changes: 11 additions & 10 deletions benchmarks/benchmarks/api/bench_knn_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,22 @@

from .. import utils


@utils.benchmark('time', timeout=60)
@utils.parametrize('k', [3, 5, 10])
@utils.parametrize('size', [50, 200, 1000])
@utils.parametrize('dim', [16, 128, 512])
def track_time(size, dim, k):
@utils.parametrize('k', [8, 64])
@utils.parametrize('size', [1000, 10000])
@utils.parametrize('dim', [4, 32, 256])
@utils.parametrize_cpu('algorithm', ['bruteforce-blas', 'bruteforce', 'kd-tree', 'nn-descent'])
@utils.parametrize_gpu('algorithm', ['bruteforce-blas', 'bruteforce', 'bruteforce-sharemem', 'nn-descent'])
def track_time(size, dim, k, algorithm):
device = utils.get_bench_device()
features = np.random.RandomState(42).randn(size, dim)
feat = torch.tensor(features, dtype=torch.float, device=device)
# dry run
for i in range(3):
dgl.knn_graph(feat, k)
for i in range(1):
dgl.knn_graph(feat, k, algorithm=algorithm)
# timing
with utils.Timer() as t:
for i in range(20):
dgl.knn_graph(feat, k)
for i in range(5):
dgl.knn_graph(feat, k, algorithm=algorithm)

return t.elapsed_secs / 20
return t.elapsed_secs / 5

0 comments on commit 1bfc311

Please sign in to comment.