def _faiss_knn(keys: torch.Tensor, queries: torch.Tensor, num_neighbors: int,
distance: str) -> Tuple[torch.Tensor, torch.Tensor]:
# https://github.com/facebookresearch/XLM/blob/master/src/model/memory/utils.py
if not is_faiss_available():
raise RuntimeError("faiss_knn requires faiss-gpu")
import faiss
assert distance in ['dot_product', 'l2']
assert keys.size(1) == queries.size(1)
metric = faiss.METRIC_INNER_PRODUCT if distance == 'dot_product' else faiss.METRIC_L2
k_ptr = _tensor_to_ptr(keys)
q_ptr = _tensor_to_ptr(queries)
scores = keys.new_zeros((queries.size(0), num_neighbors),
dtype=torch.float32)
indices = keys.new_zeros((queries.size(0), num_neighbors),
dtype=torch.int64)
s_ptr = _tensor_to_ptr(scores)
i_ptr = _tensor_to_ptr(indices)
faiss.bfKnn(FAISS_RES, metric, k_ptr, True, keys.size(0), q_ptr, True,
queries.size(0), queries.size(1), num_neighbors, s_ptr, i_ptr)
return scores, indices
def _faiss_knn(keys: torch.Tensor, queries: torch.Tensor, num_neighbors: int,
distance: str) -> Tuple[torch.Tensor, torch.Tensor]:
# https://github.com/facebookresearch/XLM/blob/master/src/model/memory/utils.py
if not is_faiss_available():
raise RuntimeError("faiss_knn requires faiss-gpu")
import faiss
metric = faiss.METRIC_INNER_PRODUCT if distance == 'dot_product' else faiss.METRIC_L2
k_ptr = _tensor_to_ptr(keys)
q_ptr = _tensor_to_ptr(queries)
scores = keys.new_zeros((queries.size(0), num_neighbors),
dtype=torch.float32)
indices = keys.new_zeros((queries.size(0), num_neighbors),
dtype=torch.int64)
s_ptr = _tensor_to_ptr(scores)
i_ptr = _tensor_to_ptr(indices)
args = faiss.GpuDistanceParams()
args.metric = metric
args.k = num_neighbors
args.dims = queries.size(1)
args.vectors = k_ptr
args.vectorsRowMajor = True
args.numVectors = keys.size(0)
args.queries = q_ptr
args.queriesRowMajor = True
args.numQueries = queries.size(0)
args.outDistances = s_ptr
args.outIndices = i_ptr
faiss.bfKnn(FAISS_RES, args)
return scores, indices
def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None,
metric=faiss.METRIC_L2):
"""search xq in xb, without building an index"""
assert xb.device == xq.device
nq, d = xq.size()
if xq.is_contiguous():
xq_row_major = True
elif xq.t().is_contiguous():
xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-)
xq_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xq_ptr = swig_ptr_from_FloatTensor(xq)
nb, d2 = xb.size()
assert d2 == d
if xb.is_contiguous():
xb_row_major = True
elif xb.t().is_contiguous():
xb = xb.t()
xb_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xb_ptr = swig_ptr_from_FloatTensor(xb)
if D is None:
D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, k)
assert D.device == xb.device
if I is None:
I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
else:
assert I.shape == (nq, k)
assert I.device == xb.device
D_ptr = swig_ptr_from_FloatTensor(D)
I_ptr = swig_ptr_from_LongTensor(I)
args = faiss.GpuDistanceParams()
args.metric = metric
args.k = k
args.dims = d
args.vectors = xb_ptr
args.vectorsRowMajor = xb_row_major
args.numVectors = nb
args.queries = xq_ptr
args.queriesRowMajor = xq_row_major
args.numQueries = nq
args.outDistances = D_ptr
args.outIndices = I_ptr
with using_stream(res):
faiss.bfKnn(res, args)
return D, I
def faiss_knn(keys: torch.Tensor, queries: torch.Tensor, num_neighbors: int,
distance: str) -> Tuple[torch.Tensor, torch.Tensor]:
""" k nearest neighbor using faiss. Users are recommended to use `k_nearest_neighbor` instead.
:param keys: tensor of (num_keys, dim)
:param queries: tensor of (num_queries, dim)
:param num_neighbors: `k`
:param distance: user can use str or faiss.METRIC_*.
:return: scores, indices in tensor
"""
if not is_faiss_available():
raise RuntimeError("_faiss_knn requires faiss-gpu")
metric_map = {
"inner_product": faiss.METRIC_INNER_PRODUCT,
"l2": faiss.METRIC_L2,
"l1": faiss.METRIC_L1,
"linf": faiss.METRIC_Linf,
"jansen_shannon": faiss.METRIC_JensenShannon
}
k_ptr = _tensor_to_ptr(keys)
q_ptr = _tensor_to_ptr(queries)
scores = keys.new_empty((queries.size(0), num_neighbors),
dtype=torch.float32)
indices = keys.new_empty((queries.size(0), num_neighbors),
dtype=torch.int64)
s_ptr = _tensor_to_ptr(scores)
i_ptr = _tensor_to_ptr(indices)
args = faiss.GpuDistanceParams()
args.metric = metric_map[distance] if isinstance(distance,
str) else distance
args.k = num_neighbors
args.dims = queries.size(1)
args.vectors = k_ptr
args.vectorsRowMajor = True
args.numVectors = keys.size(0)
args.queries = q_ptr
args.queriesRowMajor = True
args.numQueries = queries.size(0)
args.outDistances = s_ptr
args.outIndices = i_ptr
faiss.bfKnn(FAISS_RES, args)
return scores, indices
def test_dist(self):
metrics = [
faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT, faiss.METRIC_L1,
faiss.METRIC_Linf, faiss.METRIC_Canberra, faiss.METRIC_BrayCurtis,
faiss.METRIC_JensenShannon
]
for metric in metrics:
print(metric)
d = 33
k = 500
# all pairwise distance should be the same as nb = k
nb = k
nq = 20
xs = make_t(nb, d)
qs = make_t(nq, d)
res = faiss.StandardGpuResources()
# Get ground truth using IndexFlat
index = faiss.IndexFlat(d, metric)
index.add(xs)
ref_d, _ = index.search(qs, k)
out_d = np.empty((nq, k), dtype=np.float32)
# Try f32 data/queries
params = faiss.GpuDistanceParams()
params.metric = metric
params.k = -1 # all pairwise
params.dims = d
params.vectors = faiss.swig_ptr(xs)
params.numVectors = nb
params.queries = faiss.swig_ptr(qs)
params.numQueries = nq
params.outDistances = faiss.swig_ptr(out_d)
faiss.bfKnn(res, params)
# IndexFlat will sort the results, so we need to
# do the same on our end
out_d = np.sort(out_d, axis=1)
# INNER_PRODUCT is in descending order, make sure it is the same
# order
if metric == faiss.METRIC_INNER_PRODUCT:
ref_d = np.sort(ref_d, axis=1)
print('f32', np.abs(ref_d - out_d).max())
self.assertTrue(np.allclose(ref_d, out_d, atol=1e-5))
# Try float16 data/queries
xs_f16 = xs.astype(np.float16)
qs_f16 = qs.astype(np.float16)
xs_f16_f32 = xs_f16.astype(np.float32)
qs_f16_f32 = qs_f16.astype(np.float32)
index.reset()
index.add(xs_f16_f32)
ref_d_f16, _ = index.search(qs_f16_f32, k)
params.vectors = faiss.swig_ptr(xs_f16)
params.vectorType = faiss.DistanceDataType_F16
params.queries = faiss.swig_ptr(qs_f16)
params.queryType = faiss.DistanceDataType_F16
out_d_f16 = np.empty((nq, k), dtype=np.float32)
params.outDistances = faiss.swig_ptr(out_d_f16)
faiss.bfKnn(res, params)
# IndexFlat will sort the results, so we need to
# do the same on our end
out_d_f16 = np.sort(out_d_f16, axis=1)
# INNER_PRODUCT is in descending order, make sure it is the same
# order
if metric == faiss.METRIC_INNER_PRODUCT:
ref_d_f16 = np.sort(ref_d_f16, axis=1)
print('f16', np.abs(ref_d_f16 - out_d_f16).max())
self.assertTrue(np.allclose(ref_d_f16, out_d_f16, atol=4e-3))
def test_input_types(self):
d = 33
k = 5
nb = 1000
nq = 10
xs = make_t(nb, d)
qs = make_t(nq, d)
res = faiss.StandardGpuResources()
# Get ground truth using IndexFlat
index = faiss.IndexFlatL2(d)
index.add(xs)
ref_d, ref_i = index.search(qs, k)
out_d = np.empty((nq, k), dtype=np.float32)
out_i = np.empty((nq, k), dtype=np.int64)
# Try f32 data/queries, i64 out indices
params = faiss.GpuDistanceParams()
params.k = k
params.dims = d
params.vectors = faiss.swig_ptr(xs)
params.numVectors = nb
params.queries = faiss.swig_ptr(qs)
params.numQueries = nq
params.outDistances = faiss.swig_ptr(out_d)
params.outIndices = faiss.swig_ptr(out_i)
faiss.bfKnn(res, params)
self.assertTrue(np.allclose(ref_d, out_d, atol=1e-5))
self.assertGreaterEqual((out_i == ref_i).sum(), ref_i.size)
# Try int32 out indices
out_i32 = np.empty((nq, k), dtype=np.int32)
params.outIndices = faiss.swig_ptr(out_i32)
params.outIndicesType = faiss.IndicesDataType_I32
faiss.bfKnn(res, params)
self.assertEqual((out_i32 == ref_i).sum(), ref_i.size)
# Try float16 data/queries, i64 out indices
xs_f16 = xs.astype(np.float16)
qs_f16 = qs.astype(np.float16)
xs_f16_f32 = xs_f16.astype(np.float32)
qs_f16_f32 = qs_f16.astype(np.float32)
index.reset()
index.add(xs_f16_f32)
ref_d_f16, ref_i_f16 = index.search(qs_f16_f32, k)
params.vectors = faiss.swig_ptr(xs_f16)
params.vectorType = faiss.DistanceDataType_F16
params.queries = faiss.swig_ptr(qs_f16)
params.queryType = faiss.DistanceDataType_F16
out_d_f16 = np.empty((nq, k), dtype=np.float32)
out_i_f16 = np.empty((nq, k), dtype=np.int64)
params.outDistances = faiss.swig_ptr(out_d_f16)
params.outIndices = faiss.swig_ptr(out_i_f16)
params.outIndicesType = faiss.IndicesDataType_I64
faiss.bfKnn(res, params)
self.assertGreaterEqual((out_i_f16 == ref_i_f16).sum(),
ref_i_f16.size - 5)
self.assertTrue(np.allclose(ref_d_f16, out_d_f16, atol=2e-3))
def torch_replacement_knn_gpu(res,
xq,
xb,
k,
D=None,
I=None,
metric=faiss.METRIC_L2):
if type(xb) is np.ndarray:
# Forward to faiss __init__.py base method
return faiss.knn_gpu_numpy(res, xq, xb, k, D, I, metric)
nb, d = xb.size()
if xb.is_contiguous():
xb_row_major = True
elif xb.t().is_contiguous():
xb = xb.t()
xb_row_major = False
else:
raise TypeError('matrix should be row or column-major')
if xb.dtype == torch.float32:
xb_type = faiss.DistanceDataType_F32
xb_ptr = swig_ptr_from_FloatTensor(xb)
elif xb.dtype == torch.float16:
xb_type = faiss.DistanceDataType_F16
xb_ptr = swig_ptr_from_HalfTensor(xb)
else:
raise TypeError('xb must be f32 or f16')
nq, d2 = xq.size()
assert d2 == d
if xq.is_contiguous():
xq_row_major = True
elif xq.t().is_contiguous():
xq = xq.t()
xq_row_major = False
else:
raise TypeError('matrix should be row or column-major')
if xq.dtype == torch.float32:
xq_type = faiss.DistanceDataType_F32
xq_ptr = swig_ptr_from_FloatTensor(xq)
elif xq.dtype == torch.float16:
xq_type = faiss.DistanceDataType_F16
xq_ptr = swig_ptr_from_HalfTensor(xq)
else:
raise TypeError('xq must be f32 or f16')
if D is None:
D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, k)
# interface takes void*, we need to check this
assert (D.dtype == torch.float32)
if I is None:
I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
else:
assert I.shape == (nq, k)
if I.dtype == torch.int64:
I_type = faiss.IndicesDataType_I64
I_ptr = swig_ptr_from_IndicesTensor(I)
elif I.dtype == I.dtype == torch.int32:
I_type = faiss.IndicesDataType_I32
I_ptr = swig_ptr_from_IntTensor(I)
else:
raise TypeError('I must be i64 or i32')
D_ptr = swig_ptr_from_FloatTensor(D)
args = faiss.GpuDistanceParams()
args.metric = metric
args.k = k
args.dims = d
args.vectors = xb_ptr
args.vectorsRowMajor = xb_row_major
args.vectorType = xb_type
args.numVectors = nb
args.queries = xq_ptr
args.queriesRowMajor = xq_row_major
args.queryType = xq_type
args.numQueries = nq
args.outDistances = D_ptr
args.outIndices = I_ptr
args.outIndicesType = I_type
with using_stream(res):
faiss.bfKnn(res, args)
return D, I
def knn_gpu(res, xb, xq, k, D=None, I=None, metric=faiss.METRIC_L2):
"""Brute-force k-nearest neighbor on the GPU using CPU-resident numpy arrays
Supports float16 arrays and Fortran-order arrays.
"""
if xb.ndim != 2 or xq.ndim != 2:
raise TypeError('xb and xq must be matrices')
nb, d = xb.shape
nq, d2 = xq.shape
if d != d2:
raise TypeError('xq not the same dimension as xb')
if xb.flags.c_contiguous:
xb_row_major = True
elif xb.flags.f_contiguous:
xb = xb.T
xb_row_major = False
else:
raise TypeError('xb must be either C or Fortran contiguous')
if xq.flags.c_contiguous:
xq_row_major = True
elif xq.flags.f_contiguous:
xq = xq.T
xq_row_major = False
else:
raise TypeError('xq must be either C or Fortran contiguous')
if xb.dtype == np.float32 and xq.dtype == np.float32:
xb_xq_type = faiss.DistanceDataType_F32
elif xb.dtype == np.float16 and xq.dtype == np.float16:
xb_xq_type = faiss.DistanceDataType_F16
else:
raise TypeError('xb and xq must both be np.float32 or np.float16')
if D is None:
D = np.empty((nq, k), dtype=np.float32)
else:
assert D.shape == (nq, k)
assert D.dtype == np.float32
if I is None:
I = np.empty((nq, k), dtype=np.int64)
indices_type = faiss.IndicesDataType_I64
else:
assert I.shape == (nq, k)
if I.dtype == np.int64:
indices_type = faiss.IndicesDataType_I64
elif I.dtype == np.int32:
indices_type = faiss.IndicesDataType_I32
else:
raise TypeError('I must be either np.int64 or np.int32')
print('row major', xb_row_major, xq_row_major)
args = faiss.GpuDistanceParams()
args.metric = metric
args.k = k
args.dims = d
args.vectors = faiss.swig_ptr(xb)
args.vectorType = xb_xq_type
args.vectorsRowMajor = xb_row_major
args.numVectors = nb
args.queries = faiss.swig_ptr(xq)
args.queryType = xb_xq_type
args.queriesRowMajor = xq_row_major
args.numQueries = nq
args.outDistances = faiss.swig_ptr(D)
args.outIndices = faiss.swig_ptr(I)
args.outIndicesType = indices_type
faiss.bfKnn(res, args)
return D, I
def search_raw_array_pytorch(res,
xb,
xq,
k,
D=None,
I=None,
metric=faiss.METRIC_L2):
assert xb.device == xq.device
nq, d = xq.size()
if xq.is_contiguous():
xq_row_major = True
elif xq.t().is_contiguous():
xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-)
xq_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xq_ptr = swig_ptr_from_FloatTensor(xq)
nb, d2 = xb.size()
assert d2 == d
if xb.is_contiguous():
xb_row_major = True
elif xb.t().is_contiguous():
xb = xb.t()
xb_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xb_ptr = swig_ptr_from_FloatTensor(xb)
if D is None:
D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, k)
assert D.device == xb.device
if I is None:
I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
else:
assert I.shape == (nq, k)
assert I.device == xb.device
D_ptr = swig_ptr_from_FloatTensor(D)
I_ptr = swig_ptr_from_LongTensor(I)
gpu_config = faiss.GpuDistanceParams()
gpu_config.metric = metric
gpu_config.k = k
gpu_config.dims = d
gpu_config.vectors = xb_ptr
gpu_config.vectorsRowMajor = xb_row_major
gpu_config.vectorType = faiss.DistanceDataType_F32
gpu_config.numVectors = nb
gpu_config.queries = xq_ptr
gpu_config.queriesRowMajor = xq_row_major
gpu_config.queryType = faiss.DistanceDataType_F32
gpu_config.numQueries = nq
gpu_config.outDistances = D_ptr
gpu_config.outIndices = I_ptr
gpu_config.outIndicesType = faiss.DistanceDataType_F32
faiss.bfKnn(res, gpu_config)
return D, I
def search_mod_raw_array_pytorch(res,
noise_level,
xb,
xq,
k,
D=None,
I=None,
metric=faiss_mod.METRIC_L2):
assert xb.device == xq.device
nq, d = xq.size()
if xq.is_contiguous():
xq_row_major = True
elif xq.t().is_contiguous():
xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-)
xq_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xq_ptr = swig_ptr_from_FloatTensor(xq)
nb, d2 = xb.size()
assert d2 == d
if xb.is_contiguous():
xb_row_major = True
elif xb.t().is_contiguous():
xb = xb.t()
xb_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xb_ptr = swig_ptr_from_FloatTensor(xb)
if D is None:
D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, k)
assert D.device == xb.device
if I is None:
I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
else:
assert I.shape == (nq, k)
assert I.device == xb.device
D_ptr = swig_ptr_from_FloatTensor(D)
I_ptr = swig_ptr_from_LongTensor(I)
# print("xb.means()",xb.mean(1).shape,xb.mean(1))
# print("xq.means()",xq.mean(1).shape,xq.mean(1))
# print("xb.stds()",xb.std(1).shape,xb.std(1)**2)
# print("xq.stds()",xq.std(1).shape,xq.std(1)**2)
# print("xb.norms().shape",xb.norm(dim=1,p=2))
# print("xq.norms().shape",xq.norm(dim=1,p=2))
# dist,ind = wasserstein_search(xb,xq,noise_level,k)
# print("Test W Search")
# print(dist,ind)
gpu_config = faiss_mod.GpuDistanceParams()
gpu_config.metric = metric
gpu_config.k = k
gpu_config.dims = d
gpu_config.vectors = xb_ptr
gpu_config.vectorsRowMajor = xb_row_major
gpu_config.vectorType = faiss_mod.DistanceDataType_F32
gpu_config.numVectors = nb
gpu_config.queries = xq_ptr
gpu_config.queriesRowMajor = xq_row_major
gpu_config.queryType = faiss_mod.DistanceDataType_F32
gpu_config.numQueries = nq
gpu_config.outDistances = D_ptr
gpu_config.outIndices = I_ptr
gpu_config.outIndicesType = faiss_mod.DistanceDataType_F32
gpu_config.ignoreOutDistances = False
gpu_config.noise_level = 2 * noise_level**2 / xb.shape[1]
gpu_config.useWasserstein = True
faiss_mod.bfKnn(res, gpu_config)
# faiss_mod.bruteForceKnn(res, metric,
# xb_ptr, xb_row_major, nb,
# xq_ptr, xq_row_major, nq,
# d, k, D_ptr, I_ptr)
# print(D[:2],I[:2])
return D, I
