Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
knn_benchmark.rst 9.44 KB
Newer Older
Tong He's avatar
[Doc] Add KNN Benchmark Results (#3055)  
Tong He committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128 
.. _knn_benchmark:

Benchmark the performance of KNN algorithms
===========================================

In this doc, we benchmark the performance on multiple K-Nearest Neighbor algorithms implemented by :func:`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 <https://github.com/dmlc/dgl/commit/64d0f3f3554911ec06d015f1c9659180796adf9a>`_), 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.