Benchmark the performance of KNN algorithms

In this doc, we benchmark the performance on multiple K-Nearest Neighbor algorithms implemented by dgl.knn_graph().

Given a dataset of N samples with D dimensions, the common use case of KNN algorithms in graph learning is to build a KNN graph by finding the K nearest neighbors for each of the N samples among the dataset.

Empirically, the three parameters, N, D, and K, all have impact on the computation cost. To benchmark the algorithms, we pick a few represensitive datasets to cover most common scenarios:

  • A synthetic dataset with mixed gaussian samples: N = 1000, D = 3.

  • A point cloud sample from ModelNet: N = 10000, D = 3.

  • Subsets of MNIST

    • A small subset: N = 1000, D = 784

    • A medium subset: N = 10000, D = 784

    • A large subset: N = 50000, D = 784

Some notes:

  • bruteforce-sharemem is an optimized implementation of bruteforce on GPU.

  • kd-tree is currently only implemented on CPU.

  • bruteforce-blas conducts matrix multiplication, thus is memory inefficient.

  • nn-descent is an approximate algorithm, and we also report the recall rate of its result.

Results

In this section, we show the runtime and recall rate (where applicable) for the algorithms under various scenarios.

The experiments are run on an Amazon EC2 P3.2xlarge instance. This instance has 8 vCPUs with 61GB RAM, and one Tesla V100 GPU with 16GB RAM. In terms of the environment, we obtain the numbers with DGL==0.7.0(64d0f3f), PyTorch==1.8.1, CUDA==11.1 on Ubuntu 18.04.5 LTS.

  • Mixed Gaussian:

Model

CPU

GPU

K = 8

K = 64

K = 8

K = 64

bruteforce-blas

0.010

0.011

0.002

0.003

kd-tree

0.004

0.006

n/a

n/a

bruteforce

0.004

0.006

0.126

0.009

bruteforce-sharemem

n/a

n/a

0.002

0.003

nn-descent

0.014 (R: 0.985)

0.148 (R: 1.000)

0.016 (R: 0.973)

0.077 (R: 1.000)

  • Point Cloud

Model

CPU

GPU

K = 8

K = 64

K = 8

K = 64

bruteforce-blas

0.359

0.432

0.010

0.010

kd-tree

0.007

0.026

n/a

n/a

bruteforce

0.074

0.167

0.008

0.039

bruteforce-sharemem

n/a

n/a

0.004

0.017

nn-descent

0.161 (R: 0.977)

1.345 (R: 0.999)

0.086 (R: 0.966)

0.445 (R: 0.999)

  • Small MNIST

Model

CPU

GPU

K = 8

K = 64

K = 8

K = 64

bruteforce-blas

0.014

0.015

0.002

0.002

kd-tree

0.179

0.182

n/a

n/a

bruteforce

0.173

0.228

0.123

0.170

bruteforce-sharemem

n/a

n/a

0.045

0.054

nn-descent

0.060 (R: 0.878)

1.077 (R: 0.999)

0.030 (R: 0.952)

0.457 (R: 0.999)

  • Medium MNIST

Model

CPU

GPU

K = 8

K = 64

K = 8

K = 64

bruteforce-blas

0.897

0.970

0.019

0.023

kd-tree

18.902

18.928

n/a

n/a

bruteforce

14.495

17.652

2.058

2.588

bruteforce-sharemem

n/a

n/a

2.257

2.524

nn-descent

0.804 (R: 0.755)

14.108 (R: 0.999)

0.158 (R: 0.900)

1.794 (R: 0.999)

  • Large MNIST

Model

CPU

GPU

K = 8

K = 64

K = 8

K = 64

bruteforce-blas

21.829

22.135

Out of Memory

Out of Memory

kd-tree

542.688

573.379

n/a

n/a

bruteforce

373.823

432.963

10.317

12.639

bruteforce-sharemem

n/a

n/a

53.133

58.419

nn-descent

4.995 (R: 0.658)

75.487 (R: 0.999)

1.478 (R: 0.860)

15.698 (R: 0.999)

Conclusion

  • As long as you have enough memory, bruteforce-blas is the default algorithm to go with.

  • Specifically, when D is small and the data is on CPU, kd-tree is the best algorithm.