423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-12-20 22:30:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

heterogenous KNN

Browse Source 
Summary: Interface and working implementation of ragged KNN. Benchmarks (which aren't ragged) haven't slowed. New benchmark shows that ragged is faster than non-ragged of the same shape.

Reviewed By: jcjohnson

Differential Revision: D20696507

fbshipit-source-id: 21b80f71343a3475c8d3ee0ce2680f92f0fae4de
This commit is contained in:
Jeremy Reizenstein
2020-04-07 01:45:43 -07:00
committed by Facebook GitHub Bot
parent 29b9c44c0a
commit 01b5f7b228
6 changed files with 332 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
tests/test_knn.py
View File

@@ -8,6 +8,10 @@ from pytorch3d.ops.knn import _knn_points_idx_naive, knn_points_idx
class TestKNN(unittest.TestCase):
def setUp(self) -> None:
super().setUp()
torch.manual_seed(1)
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
@@ -26,7 +30,7 @@ class TestKNN(unittest.TestCase):
self.assertTrue(torch.all(idx1 == idx2))
self.assertTrue(torch.allclose(dist1, dist2))
def test_knn_vs_python_cpu(self):
def test_knn_vs_python_cpu_square(self):
""" Test CPU output vs PyTorch implementation """
device = torch.device("cpu")
Ns = [1, 4]
@@ -37,13 +41,19 @@ class TestKNN(unittest.TestCase):
sorts = [True, False]
factors = [Ns, Ds, P1s, P2s, Ks, sorts]
for N, D, P1, P2, K, sort in product(*factors):
lengths1 = torch.full((N,), P1, dtype=torch.int64, device=device)
lengths2 = torch.full((N,), P2, dtype=torch.int64, device=device)
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)
out1 = _knn_points_idx_naive(
x, y, lengths1=lengths1, lengths2=lengths2, K=K
)
out2 = knn_points_idx(
x, y, K=K, lengths1=lengths1, lengths2=lengths2, sorted=sort
)
self._check_knn_result(out1, out2, sort)
def test_knn_vs_python_cuda(self):
def test_knn_vs_python_cuda_square(self):
""" Test CUDA output vs PyTorch implementation """
device = torch.device("cuda")
Ns = [1, 4]
@@ -57,9 +67,53 @@ class TestKNN(unittest.TestCase):
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)
out1 = _knn_points_idx_naive(x, y, lengths1=None, lengths2=None, K=K)
for version in versions:
if version == 3 and K > 4:
continue
out2 = knn_points_idx(x, y, K, sort, version)
out2 = knn_points_idx(x, y, K=K, sorted=sort, version=version)
self._check_knn_result(out1, out2, sort)
def test_knn_vs_python_cpu_ragged(self):
device = torch.device("cpu")
lengths1 = torch.tensor([10, 100, 10, 100], device=device, dtype=torch.int64)
lengths2 = torch.tensor([10, 10, 100, 100], device=device, dtype=torch.int64)
N = 4
D = 3
Ks = [1, 9, 10, 11, 101]
sorts = [False, True]
factors = [Ks, sorts]
for K, sort in product(*factors):
x = torch.randn(N, lengths1.max(), D, device=device)
y = torch.randn(N, lengths2.max(), D, device=device)
out1 = _knn_points_idx_naive(
x, y, lengths1=lengths1, lengths2=lengths2, K=K
)
out2 = knn_points_idx(
x, y, lengths1=lengths1, lengths2=lengths2, K=K, sorted=sort
)
self._check_knn_result(out1, out2, sort)
def test_knn_vs_python_cuda_ragged(self):
device = torch.device("cuda")
lengths1 = torch.tensor([10, 100, 10, 100], device=device, dtype=torch.int64)
lengths2 = torch.tensor([10, 10, 100, 100], device=device, dtype=torch.int64)
N = 4
D = 3
Ks = [1, 9, 10, 11, 101]
sorts = [True, False]
versions = [0, 1, 2, 3]
factors = [Ks, sorts]
for K, sort in product(*factors):
x = torch.randn(N, lengths1.max(), D, device=device)
y = torch.randn(N, lengths2.max(), D, device=device)
out1 = _knn_points_idx_naive(
x, y, lengths1=lengths1, lengths2=lengths2, K=K
)
for version in versions:
if version == 3 and K > 4:
continue
out2 = knn_points_idx(
x, y, lengths1=lengths1, lengths2=lengths2, K=K, sorted=sort
)
self._check_knn_result(out1, out2, sort)