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)
