KNN 算法性能基准测试
本文档中，我们对 dgl.knn_graph() 实现的多种 K 近邻算法的性能进行了基准测试。
给定一个包含 N 个样本、每个样本 D 维的数据集，KNN 算法在图学习中的常见用例是通过在该数据集中为每个 N 个样本找到 K 个最近邻居来构建 KNN 图。
根据经验，N、D 和 K 这三个参数都对计算成本有影响。为了对算法进行基准测试，我们选择了一些有代表性的数据集，以涵盖大多数常见场景
混合高斯样本的合成数据集：
N = 1000,D = 3。ModelNet 中的点云样本：
N = 10000,D = 3。- MNIST 的子集
小型子集：
N = 1000,D = 784中型子集：
N = 10000,D = 784大型子集：
N = 50000,D = 784
一些注意事项
bruteforce-sharemem是bruteforce在 GPU 上的优化实现。kd-tree目前只在 CPU 上实现。bruteforce-blas执行矩阵乘法，因此内存效率较低。nn-descent是一种近似算法，我们也报告了其结果的召回率。
结果
在本节中，我们展示了算法在各种场景下的运行时间和召回率（如适用）。
实验在一台 Amazon EC2 P3.2xlarge 实例上运行。该实例配备 8 个 vCPU、61GB RAM 以及一个 16GB RAM 的 Tesla V100 GPU。环境方面，我们在安装了 DGL==0.7.0(64d0f3f)、PyTorch==1.8.1、CUDA==11.1 的 Ubuntu 18.04.5 LTS 系统上获得这些数据。
混合高斯
模型 |
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 |
不适用 |
不适用 |
bruteforce |
0.004 |
0.006 |
0.126 |
0.009 |
bruteforce-sharemem |
不适用 |
不适用 |
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) |
点云
模型 |
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 |
不适用 |
不适用 |
bruteforce |
0.074 |
0.167 |
0.008 |
0.039 |
bruteforce-sharemem |
不适用 |
不适用 |
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) |
小型 MNIST
模型 |
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 |
不适用 |
不适用 |
bruteforce |
0.173 |
0.228 |
0.123 |
0.170 |
bruteforce-sharemem |
不适用 |
不适用 |
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) |
中型 MNIST
模型 |
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 |
不适用 |
不适用 |
bruteforce |
14.495 |
17.652 |
2.058 |
2.588 |
bruteforce-sharemem |
不适用 |
不适用 |
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) |
大型 MNIST
模型 |
CPU |
GPU |
||
|---|---|---|---|---|
K = 8 |
K = 64 |
K = 8 |
K = 64 |
|
bruteforce-blas |
21.829 |
22.135 |
内存不足 |
内存不足 |
kd-tree |
542.688 |
573.379 |
不适用 |
不适用 |
bruteforce |
373.823 |
432.963 |
10.317 |
12.639 |
bruteforce-sharemem |
不适用 |
不适用 |
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) |
结论
只要您有足够的内存，
bruteforce-blas是首选的默认算法。具体来说，当
D较小且数据在 CPU 上时，kd-tree是最佳算法。