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()