unlimblue / Knn_cuda
pytorch knn [cuda version]
Stars: ✭ 86
Labels
Projects that are alternatives of or similar to Knn cuda
2016 super resolution
ICCV2015 Image Super-Resolution Using Deep Convolutional Networks
Stars: ✭ 78 (-9.3%)
Mutual labels: cuda
Torch sampling
Efficient reservoir sampling implementation for PyTorch
Stars: ✭ 68 (-20.93%)
Mutual labels: cuda
Titan
A high-performance CUDA-based physics simulation sandbox for soft robotics and reinforcement learning.
Stars: ✭ 73 (-15.12%)
Mutual labels: cuda
Pytorch Emdloss
PyTorch 1.0 implementation of the approximate Earth Mover's Distance
Stars: ✭ 82 (-4.65%)
Mutual labels: cuda
Deepjointfilter
The source code of ECCV16 'Deep Joint Image Filtering'.
Stars: ✭ 68 (-20.93%)
Mutual labels: cuda
Cuda Design Patterns
Some CUDA design patterns and a bit of template magic for CUDA
Stars: ✭ 78 (-9.3%)
Mutual labels: cuda
Mpn Cov
@ICCV2017: For exploiting second-order statistics, we propose Matrix Power Normalized Covariance pooling (MPN-COV) ConvNets, different from and outperforming those using global average pooling.
Stars: ✭ 63 (-26.74%)
Mutual labels: cuda
Parenchyma
An extensible HPC framework for CUDA, OpenCL and native CPU.
Stars: ✭ 71 (-17.44%)
Mutual labels: cuda
Mpr
Reference implementation for "Massively Parallel Rendering of Complex Closed-Form Implicit Surfaces" (SIGGRAPH 2020)
Stars: ✭ 84 (-2.33%)
Mutual labels: cuda
Modulated Deform Conv
deformable convolution 2D 3D DeformableConvolution DeformConv Modulated Pytorch CUDA
Stars: ✭ 81 (-5.81%)
Mutual labels: cuda
KNN_CUDA
- ref: kNN-CUDA
- ref: pytorch knn cuda
- author: [email protected]
Modifications
- Aten support
- pytorch v1.0+ support
- pytorch c++ extention
Performance
- dim = 5
- k = 100
- ref = 224
- query = 224
- Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
- NVIDIA GeForce 940MX
| Loop | sklearn | CUDA | Memory |
|---|---|---|---|
| 100 | 2.34 ms | 0.06 ms | 652/1024 |
| 1000 | 2.30 ms | 1.40 ms | 652/1024 |
Install
- from source
git clone https://github.com/unlimblue/KNN_CUDA.git
cd KNN_CUDA
make && make install
- from wheel
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl
And then, make sure ninja has been installed:
wget -P /usr/bin https://github.com/unlimblue/KNN_CUDA/raw/master/ninja
Usage
import torch
# Make sure your CUDA is available.
assert torch.cuda.is_available()
from knn_cuda import KNN
"""
if transpose_mode is True,
ref is Tensor [bs x nr x dim]
query is Tensor [bs x nq x dim]
return
dist is Tensor [bs x nq x k]
indx is Tensor [bs x nq x k]
else
ref is Tensor [bs x dim x nr]
query is Tensor [bs x dim x nq]
return
dist is Tensor [bs x k x nq]
indx is Tensor [bs x k x nq]
"""
knn = KNN(k=10, transpose_mode=True)
ref = torch.rand(32, 1000, 5).cuda()
query = torch.rand(32, 50, 5).cuda()
dist, indx = knn(ref, query) # 32 x 50 x 10
Note that the project description data, including the texts, logos, images, and/or trademarks,
for each open source project belongs to its rightful owner.
If you wish to add or remove any projects, please contact us at [email protected].