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)