pytorch: knn cuda在windows下的编译

前言

本篇记录一下在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包没问题。

Logo

欢迎来到由智源人工智能研究院发起的Triton中文社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。

更多推荐