def search_index_pytorch(index, x, k, D=None, I=None):
"""call the search function of an index with pytorch tensor I/O (CPU
and GPU supported)"""
assert x.is_contiguous()
n, d = x.size()
assert d == index.d
if D is None:
if x.is_cuda:
D = torch.cuda.FloatTensor(n, k)
else:
D = torch.FloatTensor(n, k)
else:
assert D.__class__ in (torch.FloatTensor, torch.cuda.FloatTensor)
assert D.size() == (n, k)
assert D.is_contiguous()
if I is None:
if x.is_cuda:
I = torch.cuda.LongTensor(n, k)
else:
I = torch.LongTensor(n, k)
else:
assert I.__class__ in (torch.LongTensor, torch.cuda.LongTensor)
assert I.size() == (n, k)
assert I.is_contiguous()
xptr = x.storage().data_ptr()
Iptr = I.storage().data_ptr()
Dptr = D.storage().data_ptr()
index.search_c(n, faiss.cast_integer_to_float_ptr(xptr),
k, faiss.cast_integer_to_float_ptr(Dptr),
faiss.cast_integer_to_long_ptr(Iptr))
return D, I
def search_index_pytorch(index, x, k, D=None, I=None):
"""call the search function of an index with pytorch tensor I/O (CPU
and GPU supported)"""
assert x.is_contiguous()
n, d = x.size()
assert d == index.d
if D is None:
if x.is_cuda:
D = torch.cuda.FloatTensor(n, k)
else:
D = torch.FloatTensor(n, k)
else:
assert D.__class__ in (torch.FloatTensor, torch.cuda.FloatTensor)
assert D.size() == (n, k)
assert D.is_contiguous()
if I is None:
if x.is_cuda:
I = torch.cuda.LongTensor(n, k)
else:
I = torch.LongTensor(n, k)
else:
assert I.__class__ in (torch.LongTensor, torch.cuda.LongTensor)
assert I.size() == (n, k)
assert I.is_contiguous()
torch.cuda.synchronize()
xptr = x.storage().data_ptr()
Iptr = I.storage().data_ptr()
Dptr = D.storage().data_ptr()
index.search_c(n, faiss.cast_integer_to_float_ptr(xptr), k,
faiss.cast_integer_to_float_ptr(Dptr),
faiss.cast_integer_to_long_ptr(Iptr))
torch.cuda.synchronize()
return D, I
def _tensor_to_ptr(input: torch.Tensor):
assert input.is_contiguous()
assert input.dtype in {torch.float32, torch.int64}
if input.dtype is torch.float32:
return faiss.cast_integer_to_float_ptr(input.storage().data_ptr() +
input.storage_offset() * 4)
else:
return faiss.cast_integer_to_long_ptr(input.storage().data_ptr() +
input.storage_offset() * 8)
def swigPtrFromTensor(x):
""" gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """
assert x.is_contiguous()
if x.dtype == torch.float32:
return faiss.cast_integer_to_float_ptr(x.storage().data_ptr() +
x.storage_offset() * 4)
if x.dtype == torch.int64:
return faiss.cast_integer_to_long_ptr(x.storage().data_ptr() +
x.storage_offset() * 8)
raise Exception("tensor type not supported: {}".format(x.dtype))
def __init__(self, linear, batch_size, topk, nprobe, npartitions, flat):
super().__init__()
linear_weight = linear.weight.detach().cpu().numpy()
self.has_bias = linear.bias is not None
if self.has_bias:
linear_bias = linear.bias.detach().cpu().numpy()
self.weights = np.column_stack([linear_weight, linear_bias])
else:
self.weights = linear_weight
if flat:
print("!! ApproxLinear: FLAT testing mode", file=sys.stderr)
self.cpu_index = faiss.index_factory(
self.weights.shape[1],
"Flat" if flat else f"IVF{npartitions},Flat",
faiss.METRIC_INNER_PRODUCT)
self.cpu_index.train(self.weights)
self.cpu_index.add(self.weights)
self.cpu_index.nprobe = nprobe
self.res = faiss.StandardGpuResources()
self.index = faiss.index_cpu_to_gpu(self.res, 0, self.cpu_index)
self.topk = topk
self.batch_size = batch_size
self.ones = torch.ones(batch_size).view(-1, 1).cuda()
self.I = torch.LongTensor(batch_size, topk).cuda()
self.D = torch.FloatTensor(batch_size, topk).cuda()
self.Dptr = faiss.cast_integer_to_float_ptr(
self.D.storage().data_ptr())
self.Iptr = faiss.cast_integer_to_long_ptr(self.I.storage().data_ptr())
self.dense = False
self.out_dim = linear_weight.shape[0]
def _knn_search(queries, data, k, return_neighbours=False, res=None):
num_queries, dim = queries.shape
if res is None:
dists, idxs = np.empty((num_queries, k), dtype=np.float32), np.empty(
(num_queries, k), dtype=np.int64)
heaps = faiss.float_maxheap_array_t()
heaps.k, heaps.nh = k, num_queries
heaps.val, heaps.ids = faiss.swig_ptr(dists), faiss.swig_ptr(idxs)
faiss.knn_L2sqr(faiss.swig_ptr(queries), faiss.swig_ptr(data), dim,
num_queries, data.shape[0], heaps)
else:
dists, idxs = torch.empty(num_queries,
k,
dtype=torch.float32,
device=queries.device), torch.empty(
num_queries,
k,
dtype=torch.int64,
device=queries.device)
faiss.bruteForceKnn(
res, faiss.METRIC_L2,
faiss.cast_integer_to_float_ptr(data.storage().data_ptr() +
data.storage_offset() * 4),
data.is_contiguous(), data.shape[0],
faiss.cast_integer_to_float_ptr(queries.storage().data_ptr() +
queries.storage_offset() * 4),
queries.is_contiguous(), num_queries, dim, k,
faiss.cast_integer_to_float_ptr(dists.storage().data_ptr() +
dists.storage_offset() * 4),
faiss.cast_integer_to_long_ptr(idxs.storage().data_ptr() +
idxs.storage_offset() * 8))
if return_neighbours:
neighbours = data[idxs.reshape(-1)].reshape(-1, k, dim)
return dists, idxs, neighbours
else:
return dists, idxs
def _swig_ptr_from_cupy_int64_array(x): # pragma: no cover
assert x.flags.c_contiguous
assert x.dtype == np.int64
data_ptr = x.__cuda_array_interface__['data'][0]
return faiss.cast_integer_to_long_ptr(data_ptr)
def swig_ptr_from_LongTensor(x):
assert x.is_contiguous()
assert x.dtype == torch.int64, 'dtype=%s' % x.dtype
return faiss.cast_integer_to_long_ptr(x.storage().data_ptr() +
x.storage_offset() * 8)
def swig_ptr_from_LongTensor(x):
""" gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """
assert x.is_contiguous()
assert x.dtype == torch.int64, 'dtype=%s' % x.dtype
return faiss.cast_integer_to_long_ptr(
x.storage().data_ptr() + x.storage_offset() * 8)
