godspeed1989 / knn_pytorch

K-Nearest Neighbor in Pytorch by CUDA

Pytorch KNN in CUDA

We calculate distance matrix and topk indices in Python. The CUDA code just gathers the nearest neighbor points with topk indices.

Install

cd knn_pytorch
make
python knn_pytorch.py

Notes

• This repository works with pytorch 1.0
• Works with batched tensors