pytorch： knn cuda编译

前言

本篇记录一下在windows系统下的pytorch版knn cuda编译。

看pytorch官方issue上写了今年可能会官方提供knn算法，但不知道啥时候会添加这个功能。远水解不了近渴，还是得上github找资源。

编译流程

编译前准备

首先把windows版knn cuda代码clone下来：

git clone --branch windows https://github.com/unlimblue/KNN_CUDA.git

这个代码是用make编译的，因此要下载chocolatey（一个类似ubuntu ape-get的东西），首先用管理员打开windows powershell，然后执行

Get-ExecutionPolicy

如果返回Restricted，则执行

Set-ExecutionPolicy Bypass -Scope Process

然后运行

Set-ExecutionPolicy Bypass -Scope Process -Force; [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.ServicePointManager]::SecurityProtocol -bor 3072; iex ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1'))

如果没有看到报错，chocolatey就下载好了，然后

choco install make

没报错make就安装好了

编译

在cmd中把自己要用的虚拟环境激活一下（没有就忽略这步）：

conda activate myenv

然后进入之前clone的KNN_CUDA根目录下，运行

make
make install

安装完后，查看一下环境

pip list

看到KNN_CUDA包名，就表明装好了。

测试（可选）

如果想看下自己编译的包有没有问题，可以用pytest：

pip install pytest
pytest tests/test_knn_cuda.py

然后就会进行大概三分钟的测试，最后测试例全pass，就表明编出来的KNN_CUDA包没问题。