def test_nn_cuda_error(self):
"""
Check that nn_points_idx throws an error if cpu tensors
are given as input.
"""
x = torch.randn(1, 1, 3)
y = torch.randn(1, 1, 3)
with self.assertRaises(Exception) as err:
_C.nn_points_idx(x, y)
self.assertTrue("Not implemented on the CPU" in str(err.exception))
def _test_nn_helper(self, device):
for D in [3, 4]:
for N in [1, 4]:
for P1 in [1, 8, 64, 128]:
for P2 in [32, 128]:
x = torch.randn(N, P1, D, device=device)
y = torch.randn(N, P2, D, device=device)
# _C.nn_points_idx should dispatch
# to the cpp or cuda versions of the function
# depending on the input type.
idx1 = _C.nn_points_idx(x, y)
idx2 = TestNearestNeighborPoints.nn_points_idx_naive(
x, y)
self.assertTrue(idx1.size(1) == P1)
self.assertTrue(torch.all(idx1 == idx2))
def test_nn_cuda(self):
"""
Test cuda output vs naive python implementation.
"""
device = torch.device("cuda:0")
for D in [3, 4]:
for N in [1, 4]:
for P1 in [1, 8, 64, 128]:
for P2 in [32, 128]:
x = torch.randn(N, P1, D, device=device)
y = torch.randn(N, P2, D, device=device)
# _C.nn_points_idx should dispatch
# to the cpp or cuda versions of the function
# depending on the input type.
idx1 = _C.nn_points_idx(x, y)
idx2 = TestNearestNeighborPoints.nn_points_idx_naive(
x, y)
self.assertTrue(idx1.size(1) == P1)
self.assertTrue(torch.all(idx1 == idx2))
def nn_points_idx(p1, p2, p2_normals=None) -> torch.Tensor:
"""
Compute the coordinates of nearest neighbors in pointcloud p2 to points in p1.
Args:
p1: FloatTensor of shape (N, P1, D) giving a batch of pointclouds each
containing P1 points of dimension D.
p2: FloatTensor of shape (N, P2, D) giving a batch of pointclouds each
containing P2 points of dimension D.
p2_normals: [optional] FloatTensor of shape (N, P2, D) giving
normals for p2. Default: None.
Returns:
3-element tuple containing
- **p1_nn_points**: FloatTensor of shape (N, P1, D) where
p1_neighbors[n, i] is the point in p2[n] which is
the nearest neighbor to p1[n, i].
- **p1_nn_idx**: LongTensor of shape (N, P1) giving the indices of
the neighbors.
- **p1_nn_normals**: Normal vectors for each point in p1_neighbors;
only returned if p2_normals is passed
else return [].
"""
N, P1, D = p1.shape
with torch.no_grad():
p1_nn_idx = _C.nn_points_idx(
p1.contiguous(), p2.contiguous()
) # (N, P1)
p1_nn_idx_expanded = p1_nn_idx.view(N, P1, 1).expand(N, P1, D)
p1_nn_points = p2.gather(1, p1_nn_idx_expanded)
if p2_normals is None:
p1_nn_normals = []
else:
if p2_normals.shape != p2.shape:
raise ValueError("p2_normals has incorrect shape.")
p1_nn_normals = p2_normals.gather(1, p1_nn_idx_expanded)
return p1_nn_points, p1_nn_idx, p1_nn_normals
def nn_cpp():
_C.nn_points_idx(x.contiguous(), y.contiguous())
torch.cuda.synchronize()
def nn_cpu():
_C.nn_points_idx(x.contiguous(), y.contiguous())
def knn():
_C.nn_points_idx(x, y)
def knn():
_C.nn_points_idx(x, y)
torch.cuda.synchronize()
