def build_edges(spatial, r_max, k_max, return_indices=False, target_spatial=None): if k_max > 200: if device == "cuda": res = faiss.StandardGpuResources() if target_spatial is None: D, I = faiss.knn_gpu(res, spatial, spatial, k_max) else: D, I = faiss.knn_gpu(res, spatial, target_spatial, k_max) elif device == "cpu": index = faiss.IndexFlatL2(spatial.shape[1]) index.add(spatial) if target_spatial is None: D, I = index.search(spatial, k_max) else: D, I = index.search(target_spatial, k_max) else: if target_spatial is None: knn_object = ops.knn_points(spatial.unsqueeze(0), spatial.unsqueeze(0), K=k_max, return_sorted=False) else: knn_object = ops.knn_points( spatial.unsqueeze(0), target_spatial.unsqueeze(0), K=k_max, return_sorted=False, ) I = knn_object.idx[0] D = knn_object.dists[0] # Overlay the "source" hit ID onto each neighbour ID (this is necessary as the FAISS algo does some shortcuts) ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1).T edge_list = torch.stack([ind[D <= r_max**2], I[D <= r_max**2]]) # Remove self-loops edge_list = edge_list[:, edge_list[0] != edge_list[1]] if return_indices: return edge_list, D, I, ind else: return edge_list
def build_edges(query, database, indices=None, r_max=1.0, k_max=10, return_indices=False): """ NOTE: These KNN/FRNN algorithms return the distances**2. Therefore we need to be careful when comparing them to the target distances (r_val, r_test), and to the margin parameter (which is L1 distance) """ if FRNN_AVAILABLE: Dsq, I, nn, grid = frnn.frnn_grid_points( points1=query.unsqueeze(0), points2=database.unsqueeze(0), lengths1=None, lengths2=None, K=k_max, r=r_max, grid=None, return_nn=False, return_sorted=True, ) I = I.squeeze().int() ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1).T.int() positive_idxs = I >= 0 edge_list = torch.stack([ind[positive_idxs], I[positive_idxs]]).long() else: if device == "cuda": res = faiss.StandardGpuResources() Dsq, I = faiss.knn_gpu(res=res, xq=query, xb=database, k=k_max) elif device == "cpu": index = faiss.IndexFlatL2(database.shape[1]) index.add(database) Dsq, I = index.search(query, k_max) ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1).T.int() edge_list = torch.stack([ind[Dsq <= r_max**2], I[Dsq <= r_max**2]]) # Reset indices subset to correct global index if indices is not None: edge_list[0] = indices[edge_list[0]] # Remove self-loops edge_list = edge_list[:, edge_list[0] != edge_list[1]] if return_indices: return edge_list, Dsq, I, ind else: return edge_list
def test_knn_gpu_datatypes(self): torch.manual_seed(10) d = 10 nb = 1024 nq = 5 k = 10 res = faiss.StandardGpuResources() # make GT on torch cpu and test using IndexFlatL2 xb = torch.rand(nb, d, dtype=torch.float32) xq = torch.rand(nq, d, dtype=torch.float32) index = faiss.IndexFlatL2(d) index.add(xb) gt_D, gt_I = index.search(xq, k) xb_c = xb.cuda().half() xq_c = xq.cuda().half() # use i32 output indices D = torch.zeros(nq, k, device=xb_c.device, dtype=torch.float32) I = torch.zeros(nq, k, device=xb_c.device, dtype=torch.int32) faiss.knn_gpu(res, xb_c, xq_c, k, D, I) self.assertTrue(torch.equal(I.long().cpu(), gt_I)) self.assertLess((D.float().cpu() - gt_D).abs().max(), 1.5e-3) # Test using numpy D = np.zeros((nq, k), dtype=np.float32) I = np.zeros((nq, k), dtype=np.int32) xb_c = xb.half().numpy() xq_c = xq.half().numpy() faiss.knn_gpu(res, xb_c, xq_c, k, D, I) self.assertTrue(torch.equal(torch.from_numpy(I).long(), gt_I)) self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1.5e-3)
def _knn_search(queries, data, k): """ Perform exact knn search (should be replaced with approximate) Return the k nearest keys """ if torch.cuda.is_available( ): # not the best way but should let me know that gpu is being used res = faiss.StandardGpuResources() D, I = faiss.knn_gpu(res, queries, data, k) return D.detach().cpu().numpy(), I.detach().cpu().numpy() queries, data = queries.detach().numpy(), data.detach().numpy() return faiss.knn(queries, data, k) #(distances, indices)
def build_knn(spatial, k): if device == "cuda": res = faiss.StandardGpuResources() _, I = faiss.knn_gpu(res, spatial, spatial, k_max) elif device == "cpu": index = faiss.IndexFlatL2(spatial.shape[1]) index.add(spatial) _, I = index.search(spatial, k_max) I = I[:, 1:] ind = torch.Tensor.repeat_interleave( torch.arange(I.shape[0], device=device), I.shape[1]) edge_list = torch.stack([ind, I]).shape return edge_list
def build_knn(spatial, k): if device == "cuda": res = faiss.StandardGpuResources() _, I = faiss.knn_gpu(res, spatial, spatial, k_max) elif device == "cpu": index = faiss.IndexFlatL2(spatial.shape[1]) index.add(spatial) _, I = index.search(spatial, k_max) ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1).T edge_list = torch.stack([ind, I]) # Remove self-loops edge_list = edge_list[:, edge_list[0] != edge_list[1]] return edge_list
def build_edges(spatial, r_max, k_max, return_indices=False): if device == "cuda": res = faiss.StandardGpuResources() D, I = faiss.knn_gpu(res, spatial, spatial, k_max) elif device == "cpu": index = faiss.IndexFlatL2(spatial.shape[1]) index.add(spatial) D, I = index.search(spatial, k_max) ind = torch.Tensor.repeat(I[:, 0], (I.shape[1] - 1, 1), 1).T D, I = D[:, 1:], I[:, 1:] edge_list = torch.stack([ind[D <= r_max**2], I[D <= r_max**2]]) if return_indices: return edge_list, D, I, ind else: return edge_list
def build_knn(query, database, k): if device == "cuda": res = faiss.StandardGpuResources() _, I = faiss.knn_gpu(res=res, xq=query, xb=database, k=k) elif device == "cpu": index = faiss.IndexFlatL2(database.shape[1]) index.add(query) _, I = index.search(query, k) ind = torch.Tensor.repeat( torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1 ).T edge_list = torch.stack([ind, I]) # Remove self-loops edge_list = edge_list[:, edge_list[0] != edge_list[1]] return edge_list
def build_edges(spatial, r_max, k_max, return_indices=False): if device == "cuda": res = faiss.StandardGpuResources() D, I = faiss.knn_gpu(res, spatial, spatial, k_max) elif device == "cpu": index = faiss.IndexFlatL2(spatial.shape[1]) index.add(spatial) D, I = index.search(spatial, k_max) # Overlay the "source" hit ID onto each neighbour ID (this is necessary as the FAISS algo does some shortcuts) ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device), (I.shape[1], 1), 1).T edge_list = torch.stack([ind[D <= r_max**2], I[D <= r_max**2]]) # Remove self-loops edge_list = edge_list[:, edge_list[0] != edge_list[1]] if return_indices: return edge_list, D, I, ind else: return edge_list
def test_knn_gpu(self): torch.manual_seed(10) d = 32 nb = 1024 nq = 10 k = 10 res = faiss.StandardGpuResources() # make GT on torch cpu and test using IndexFlatL2 xb = torch.rand(nb, d, dtype=torch.float32) xq = torch.rand(nq, d, dtype=torch.float32) index = faiss.IndexFlatL2(d) index.add(xb) gt_D, gt_I = index.search(xq, k) # for the GPU, we'll use a non-default stream s = torch.cuda.Stream() with torch.cuda.stream(s): # test numpy inputs xb_np = xb.numpy() xq_np = xq.numpy() for xq_row_major in True, False: for xb_row_major in True, False: if not xq_row_major: xq_c = to_column_major_numpy(xq_np) assert not xq_c.flags.contiguous else: xq_c = xq_np if not xb_row_major: xb_c = to_column_major_numpy(xb_np) assert not xb_c.flags.contiguous else: xb_c = xb_np D, I = faiss.knn_gpu(res, xb_c, xq_c, k) self.assertTrue(torch.equal(torch.from_numpy(I), gt_I)) self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1e-4) # test torch (cpu, gpu) inputs for is_cuda in True, False: for xq_row_major in True, False: for xb_row_major in True, False: if is_cuda: xq_c = xq.cuda() xb_c = xb.cuda() else: # also test torch cpu tensors xq_c = xq xb_c = xb if not xq_row_major: xq_c = to_column_major_torch(xq) assert not xq_c.is_contiguous() if not xb_row_major: xb_c = to_column_major_torch(xb) assert not xb_c.is_contiguous() D, I = faiss.knn_gpu(res, xb_c, xq_c, k) self.assertTrue(torch.equal(I.cpu(), gt_I)) self.assertLess((D.cpu() - gt_D).abs().max(), 1e-4) # test on subset try: # This internally uses the current pytorch stream D, I = faiss.knn_gpu(res, xb_c, xq_c[6:8], k) except TypeError: if not xq_row_major: # then it is expected continue # otherwise it is an error raise self.assertTrue(torch.equal(I.cpu(), gt_I[6:8])) self.assertLess((D.cpu() - gt_D[6:8]).abs().max(), 1e-4)