mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
870290df34
Summary: Implements K-Nearest Neighbors with C++ and CUDA versions. KNN in CUDA is highly nontrivial. I've implemented a few different versions of the kernel, and we heuristically dispatch to different kernels based on the problem size. Some of the kernels rely on template specialization on either D or K, so we use template metaprogramming to compile specialized versions for ranges of D and K. These kernels are up to 3x faster than our existing 1-nearest-neighbor kernels, so we should also consider swapping out `nn_points_idx` to use these kernels in the backend. I've been working mostly on the CUDA kernels, and haven't converged on the correct Python API. I still want to benchmark against FAISS to see how far away we are from their performance. Reviewed By: bottler Differential Revision: D19729286 fbshipit-source-id: 608ffbb7030c21fe4008f330522f4890f0c3c21a
66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
|
|
import unittest
|
|
from itertools import product
|
|
import torch
|
|
|
|
from pytorch3d.ops.knn import _knn_points_idx_naive, knn_points_idx
|
|
|
|
|
|
class TestKNN(unittest.TestCase):
|
|
def _check_knn_result(self, out1, out2, sorted):
|
|
# When sorted=True, points should be sorted by distance and should
|
|
# match between implementations. When sorted=False we we only want to
|
|
# check that we got the same set of indices, so we sort the indices by
|
|
# index value.
|
|
idx1, dist1 = out1
|
|
idx2, dist2 = out2
|
|
if not sorted:
|
|
idx1 = idx1.sort(dim=2).values
|
|
idx2 = idx2.sort(dim=2).values
|
|
dist1 = dist1.sort(dim=2).values
|
|
dist2 = dist2.sort(dim=2).values
|
|
if not torch.all(idx1 == idx2):
|
|
print(idx1)
|
|
print(idx2)
|
|
self.assertTrue(torch.all(idx1 == idx2))
|
|
self.assertTrue(torch.allclose(dist1, dist2))
|
|
|
|
def test_knn_vs_python_cpu(self):
|
|
""" Test CPU output vs PyTorch implementation """
|
|
device = torch.device('cpu')
|
|
Ns = [1, 4]
|
|
Ds = [2, 3]
|
|
P1s = [1, 10, 101]
|
|
P2s = [10, 101]
|
|
Ks = [1, 3, 10]
|
|
sorts = [True, False]
|
|
factors = [Ns, Ds, P1s, P2s, Ks, sorts]
|
|
for N, D, P1, P2, K, sort in product(*factors):
|
|
x = torch.randn(N, P1, D, device=device)
|
|
y = torch.randn(N, P2, D, device=device)
|
|
out1 = _knn_points_idx_naive(x, y, K, sort)
|
|
out2 = knn_points_idx(x, y, K, sort)
|
|
self._check_knn_result(out1, out2, sort)
|
|
|
|
def test_knn_vs_python_cuda(self):
|
|
""" Test CUDA output vs PyTorch implementation """
|
|
device = torch.device('cuda')
|
|
Ns = [1, 4]
|
|
Ds = [2, 3, 8]
|
|
P1s = [1, 8, 64, 128, 1001]
|
|
P2s = [32, 128, 513]
|
|
Ks = [1, 3, 10]
|
|
sorts = [True, False]
|
|
versions = [0, 1, 2, 3]
|
|
factors = [Ns, Ds, P1s, P2s, Ks, sorts]
|
|
for N, D, P1, P2, K, sort in product(*factors):
|
|
x = torch.randn(N, P1, D, device=device)
|
|
y = torch.randn(N, P2, D, device=device)
|
|
out1 = _knn_points_idx_naive(x, y, K, sorted=sort)
|
|
for version in versions:
|
|
if version == 3 and K > 4:
|
|
continue
|
|
out2 = knn_points_idx(x, y, K, sort, version)
|
|
self._check_knn_result(out1, out2, sort)
|