什么是KNN_cuda

​KNN_CUDA 是一个用 CUDA 实现的计算 k 近邻搜索的项目。

​KNN 搜索是指通过计算 query 的向量和 reference 集合内的向量的相似度,并找出最相似的 k 个。它是许多研究和工业领域的一个问题,如三维物体渲染,基于内容的图像检索,统计(熵和分歧估计),生物学(基因分类)等。

​一般的解决方法是串行的算出相似度,然后用容量为 k 的最小堆(或最大堆)来找找出最相似的 k 个。但这种传统的计算方法往往耗时很长。让人们难以接受。而每两个向量之间的相似度计算是独立的,这就让这个问题可以很好地被并行的实现。

​而KNN_CUDA项目通过 NVIDA 的 API 在显卡上实现并行的求解 KNN 问题。

​我用KNN_CUDA的原因是在复现point-mae[1]的源码[2]的时候,其在对点云数据的处理过程中,用到了KNN的算法,而点云的点数量比较大,采用KNN_CUDA库则能很好的降低运算时间。

遇到的问题

在照着point-mae库中附带的readme文件安装对应的环境后,每次在执行过程中,都会在ninja的部分报编译的错误,具体的错误我没有记录(只记得有一个“Command ['ninja', '-v'] returned non-zero exit status 1”的错误),我照着一些帖子的方法依次修改,甚至修改了库中的头文件,但是还是有新的错误报出来,所以关键的问题并不在这里。

​在不断的尝试中,每次都会有ImportError: No module named 'knn'的错误报出来,这是在import knn_cuda时报出的错误,所以当时的我认为是安装没有到位,于是卸载了这个库,再按readme中的步骤,重新安装了一遍,但是错误没有消除。在这个时候,我才意识到,我安装的一直都是 linux版本 的knn_cuda(knn_cuda的原生版本就是linux上的,同时linux的安装也十分方便)。于是我将knn_cuda再次删除,结合knn_cuda源码库[3]和一篇知乎上的帖子[4]中的方法以make的方式重新安装了一遍windows版本的knn_cuda。但是还是报了ImportError的错误,而且在make的环节就报了这个错误,更加离谱。

​之后我反复的搜索,在knn_cuda的github的issue页面反复的寻找,也没有找到解决的方法。这个时候我万念俱灰,在思考换个linux系统的服务器实验或者放弃。

没有想到装完之后还会遇到问题,我之前在使用的时候就遇到过一直卡在from knn_cuda import KNN的经历,当时的我以为又是环境的问题,于是有重新安装了一遍环境,解决了这个问题,但是之后又遇到了这个问题,我意识到这个问题是可以复现的,所以不应该是环境的问题,于是打算解决这个问题。当然,网络上还是没有直接的案例,还是需要多方搜索和自己的摸索。

(写在开头: 新建一个干净的conda环境,然后重新安装相应版本的knn_cuda

​后面我重新思考了一下,既然knn_cuda的readme文件中已经给出了windows下的安装步骤,而且是在2021年新加的代码,这就说明windows系统下使用knn_cuda是可行的,但我之前从没有遇到过安装完库还在报ImportError的情况。这超出了我的经验范围,也没有可供借鉴的情况,于是我选择在conda中新建一个环境重新装一遍这个库。

​换了个新的环境后,直接解决了问题,非常流畅地在windows下安装了knn_cuda,同时也成功运行了pointmae的源码。

​难以置信。

​后续我复盘了一下,应该是按照pointmae的readme逐步安装的时候,按照这个命令(pip install "git+https://github.com/erikwijmans/Pointnet2_PyTorch.git#egg=pointnet2_ops&subdirectory=pointnet2_ops_lib")误安装了linux版本的knn_cuda,然后用pip uninstall没有删除干净,这就导致在后续安装windows版本的knn_cuda的时候在环境中产生了冲突。

(写在开头: 找到你Anaconda3\envs\(你的环境名)\Lib\site-packages\knn_cuda\csrc\_ext\knn中的一个lock的文件,删掉即可

解决方法就是通过debug一步步确定问题所在,最直观的就是代码一直卡在from knn_cuda import KNN的语句,然后打断点依次进入,最后发现是在cpp_extension.py的文件中的if baton.try_acquire():语句,在这个语句中,他会直接进入else的分支,最后一直卡在baton.wait()语句上,看了网络的帖子发现这个一般是资源被锁了,于是我跟着self中的path找到了位于knn_cuda文件夹中的lock文件,删除后,故障排除,程序再也没有卡在import的语句上了。

​有时候复杂的问题只需要简单的处理步骤。

​确实,一般来说有着有模有样的readme指导文件和教程的安装说明,意味着这个软件的安装多半就是可复现的,一些共性的问题在论坛上也会有些反馈。在这种情况下遇到一个没有解决的问题的时候(这个问题在github的issue上有反映,但是没有人给出解决方案),多半是环境起了冲突,试着重装一下环境可能是最好的选择。

同样的,如果一个错误遇到两次了,那么一定是一个可复现的错误,那肯定是能够解决的。

​学习的路是无止境的,多学点什么,总不会亏待自己。

[1] Pang, Yatian, Wenxiao Wang, Francis E. H. Tay, W. Liu, Yonghong Tian and Liuliang Yuan. “Masked Autoencoders for Point Cloud Self-supervised Learning.” European Conference on Computer Vision (2022).

[2] 2022. "Point-MAE". [Online]. Available: https://github.com/Pang-Yatian/Point-MAE .

[3] 2021. "KNN_CUDA". [Online]. Available: https://github.com/unlimblue/KNN_CUDA .

[4] 2022. "【Win11】在Windows下编译安装KNN_CUDA". [Online]. Available: https://zhuanlan.zhihu.com/p/582688478 .