def knn_faiss(database, queries, dim=300, k=5): """Get KNNs of all vectors in queries. Args: database - array of vectors from which neighbors are to be searched queries - array of vectors for which neighbors are to be searched dim - dimension of vectors k - k in KNN returns: dist - A matrix of shape (queries.shape[0], k) - distances of it's KNNs of each query idxs - A matrix of shape (queries.shape[0], k) - indicies of it's KNNs of each query """ database = database.astype('float32') queries = queries.astype('float32') index = faiss.IndexFlatL2(dim) gpu_index = faiss.index_cpu_to_gpu(res, 0, index) gpu_index.add(database) dist, idxs = gpu_index.search(queries, k) return dist, idxs
def search_against_fragment(fragment: NpArray, test_vectors: NpArray) -> TwoNpArrays: # build a flat index (CPU) index_flat = faiss.IndexFlatL2(d) # make it into a GPU index gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) print("loading", fragment) landmark_data = np.load(fragment) index_names, index_vectors = landmark_data["images"], landmark_data[ "features"] print("index_names:", index_names.shape) print("vectors shape", index_vectors.shape) gpu_index_flat.add(index_vectors) print("total size of index:", gpu_index_flat.ntotal) # print("sanity search...") # distances, index = gpu_index_flat.search(index_vectors[:10], K) # actual search # print(index[0]) # print(distances[0]) print("searching") distances, index = gpu_index_flat.search(test_vectors, K) # actual search index = index_names[index] print(index[:10, :5]) print(distances[:10, :5]) return index, distances
def load_globally(word_vectors_fpath: str, faiss_gpu: bool): global wv global index_faiss print("Loading word vectors from:", word_vectors_fpath) tic = time() if word_vectors_fpath.endswith(".vec.gz"): wv = KeyedVectors.load_word2vec_format(word_vectors_fpath, binary=False, unicode_errors="ignore") else: wv = KeyedVectors.load(word_vectors_fpath) print("Loaded in {} sec.".format(time() - tic)) wv.init_sims(replace=True) if faiss_gpu: res = faiss.StandardGpuResources() # use a single GPU index_flat = faiss.IndexFlatIP( wv.vector_size) # build a flat (CPU) index index_faiss = faiss.index_cpu_to_gpu( res, GPU_DEVICE, index_flat) # make it into a gpu index index_faiss.add(wv.syn0norm) # add vectors to the index else: index_faiss = faiss.IndexFlatIP(wv.vector_size) index_faiss.add(wv.syn0norm) return wv
def __init__(self, target, nprobe=128, num_gpu=None, index_factory_str=None, verbose=False, mode='proxy', using_gpu=True): self._res_list = [] found_gpu = len(os.environ['CUDA_VISIBLE_DEVICES'].split(",")) if found_gpu == 0: raise RuntimeError( "No GPU found. Please export CUDA_VISIBLE_DEVICES") if num_gpu is None or num_gpu > found_gpu: num_gpu = found_gpu print('[faiss gpu] #GPU: {}'.format(num_gpu)) size, dim = target.shape assert size > 0, "size: {}".format(size) index_factory_str = "IVF{},PQ{}".format( min(8192, 16 * round(np.sqrt(size))), 32) if index_factory_str is None else index_factory_str cpu_index = faiss.index_factory(dim, index_factory_str) cpu_index.nprobe = nprobe if mode == 'proxy': co = faiss.GpuClonerOptions() co.useFloat16 = True co.usePrecomputed = False index = faiss.IndexProxy() for i in range(num_gpu): res = faiss.StandardGpuResources() self._res_list.append(res) sub_index = faiss.index_cpu_to_gpu( res, i, cpu_index, co) if using_gpu else cpu_index index.addIndex(sub_index) elif mode == 'shard': raise NotImplementedError else: raise KeyError("Unknown index mode") index = faiss.IndexIDMap(index) index.verbose = verbose # get nlist to decide how many samples used for training nlist = int([ item for item in index_factory_str.split(",") if 'IVF' in item ][0].replace("IVF", "")) # training if not index.is_trained: indexes_sample_for_train = np.random.randint(0, size, nlist * 256) index.train(target[indexes_sample_for_train]) # add with ids target_ids = np.arange(0, size) index.add_with_ids(target, target_ids) self.index = index
def graph_init(self, X): d = X.shape[1] res = faiss.StandardGpuResources() index_flat = faiss.IndexFlatL2(d) gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) gpu_index_flat.add(X) N = X.shape[0] c = time.time() self.D, self.I = gpu_index_flat.search(X, self.m) elapsed = time.time() - c rr = np.zeros(N); A = np.zeros((N, N)); for i in range(N): di = np.array(self.D[i, 1:self.k + 2]) if(i==0): print(di) rr[i] = 0.5 * (self.k * di[self.k] - sum(di[0:self.k])) id = self.I[i, 1:self.k + 2]; A[i, id] = (di[self.k] - di) / (self.k * di[self.k] - sum(di[1:self.k]) + 2.2204e-16); self.r = sum(rr) / len(rr); self.Lambda = self.r; A0 = (A + A.transpose()) / 2; D0 = np.eye(N) * A0.sum(axis=1); L0 = D0 - A0; return L0,A0
def doRetrieval(Q, X, k=100, verbose=True): res = faiss.StandardGpuResources() if verbose: print("creating indexFlatl2") index = faiss.IndexFlatL2(X.shape[1]) if verbose: print("put to gpu") index = faiss.index_cpu_to_gpu(res, 0, index) if verbose: print("adding index to faiss") # X shape: nxd # split into 2 chunks index.add(X) if verbose: print("num of index: " + str(index.ntotal)) if verbose: print("searching") start = time.time() D, I = index.search(Q, k) if verbose: print('Computing dot product') elapse = time.time() - start if verbose: print(elapse) return D, I
def __init__(self, dim=10, nlist=100, gpu=-1): self.dim = dim self.nlist = nlist #聚类中心的个数 #self.index = faiss.IndexFlatL2(dim) # build the index quantizer = faiss.IndexFlatL2(dim) # the other index # faiss.METRIC_L2: faiss定义了两种衡量相似度的方法(metrics), # 分别为faiss.METRIC_L2 欧式距离、 faiss.METRIC_INNER_PRODUCT 向量内积 # here we specify METRIC_L2, by default it performs inner-product search self.index = faiss.IndexIVFFlat(quantizer, dim, self.nlist, faiss.METRIC_L2) try: if gpu >= 0: if gpu == 0: # use a single GPU res = faiss.StandardGpuResources() gpu_index = faiss.index_cpu_to_gpu(res, 0, self.index) else: gpu_index = faiss.index_cpu_to_all_gpus(self.index) self.index = gpu_index except: pass # data self.xb = None
def KNN(query, gallery, K=10, mode='ones'): '''retrieves the K-Nearest Neighbors in the gallery''' d = query.shape[1] query = L2norm(query) gallery = L2norm(gallery) res = faiss.StandardGpuResources() index_flat = faiss.IndexFlatL2(d) if torch.cuda.is_available(): gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) gpu_index_flat.add(gallery) D, I = gpu_index_flat.search(query, K) else: index_flat.add(gallery) D, I = index_flat.search(query, K) if mode == 'lin': weights = (float(K) - np.arange(0, K)) / float(K) elif mode == 'exp': weights = np.exp(-np.arange(0, K)) elif mode == 'ones': weights = np.ones(K) weights_sum = weights.sum() new_queries = [] for i in range(len(query)): idx = I[i, :K] to_consider = gallery[idx, :] new_queries.append(np.dot(weights, to_consider) / weights_sum) new_queries = np.asarray(new_queries, dtype=np.float32) return new_queries
def detect_border_gpu(pcs): dim = 3 sem_all = np.array( list(chain.from_iterable([[i] * len(pc) for i, pc in enumerate(pcs)]))) pcs_all = np.concatenate(pcs, axis=0).astype(np.float32) index = np.random.choice(len(pcs_all), int(len(pcs_all) / 10), replace=False) pcs_chosen = pcs_all[index] sem_chosen = sem_all[index] res = faiss.StandardGpuResources() index = faiss.IndexFlatL2(dim) gpu_index = faiss.index_cpu_to_gpu(res, 0, index) gpu_index.add(pcs_chosen) k = 1024 nq = 100000 thresh = 0.01 border = [] for j in range(int(len(pcs_all) / nq) + 1): query = pcs_all[j * nq:(j + 1) * nq] print("{} queries for {}th batch".format(len(query), j)) D, I = gpu_index.search(query, k) for i, (dis, ind) in enumerate(zip(D, I)): neighbor = ind[dis < thresh] if len(np.unique(sem_chosen[neighbor])) > 1: border.append(j * nq + i) return np.array(border)
def __init__(self, vectors, identifiers, dim=512, gpu=True, inbuilt_index=False): """ :param vectors: :param identifiers: :param dim: :param gpu: :param inbuilt_index: """ Embeddings.validate(vectors, identifiers) self.__dimension = dim self.__vectors = vectors self.__identifiers = identifiers self.__gpu = gpu self.__inbuilt_index = inbuilt_index self.__quantizer = faiss.IndexFlatL2(dim) # the other index if self.__inbuilt_index: self.__indexmap = faiss.IndexIDMap2(self.__quantizer) else: self.__indexmap = self.__quantizer if self.__gpu: self.__index = faiss.index_cpu_to_gpu(GPU, 0, self.__indexmap) else: self.__index = self.__indexmap self.__add()
def fit(self, X): res = faiss.StandardGpuResources() self.index = faiss.IndexFlatL2(X.shape[1]) if self.device == 'gpu': self.index = faiss.index_cpu_to_gpu(res, 0, self.index) self.index.add(X.astype(np.float32))
def train_coarse_quantizer(data, quantizer_path, num_clusters, hnsw=False, niter=10, cuda=False): d = data.shape[1] index_flat = faiss.IndexFlatL2(d) # make it into a gpu index if cuda: res = faiss.StandardGpuResources() index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) clus = faiss.Clustering(d, num_clusters) clus.verbose = True clus.niter = niter clus.train(data, index_flat) centroids = faiss.vector_float_to_array(clus.centroids) centroids = centroids.reshape(num_clusters, d) if hnsw: quantizer = faiss.IndexHNSWFlat(d, 32) quantizer.hnsw.efSearch = 128 quantizer.train(centroids) quantizer.add(centroids) else: quantizer = faiss.IndexFlatL2(d) quantizer.add(centroids) faiss.write_index(quantizer, quantizer_path)
def search_against_fragment(train_features: np.ndarray, test_features: np.ndarray) \ -> Tuple[np.ndarray, np.ndarray]: if USE_GPU: # build a flat index (CPU) if USE_COSINE_DIST: index_flat = faiss.IndexFlat(DIMS, faiss.METRIC_INNER_PRODUCT) else: index_flat = faiss.IndexFlatL2(DIMS) # make it into a GPU index index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) else: index_flat = faiss.IndexFlatIP(DIMS) index_flat.add(train_features) print("total size of the database:", index_flat.ntotal) # print("sanity search...") # distances, index = index_flat.search(train_features[:10], K) # actual search # print(index[:10]) # print(distances[:10]) print("searching") distances, index = index_flat.search(test_features, K) # actual search dprint(index) dprint(distances) dprint(describe(distances.flatten())) return index, distances
def build_faiss_index(nd_feats_array, mode): """ build index on multi GPUs :param nd_feats_array: :param mode: 0: CPU; 1: GPU; 2: Multi-GPU :return: """ d = nd_feats_array.shape[1] cpu_index = faiss.IndexFlatL2(d) # build the index on CPU if mode == 0: print("[INFO] Is trained? >> {}".format(cpu_index.is_trained)) cpu_index.add(nd_feats_array) # add vectors to the index print("[INFO] Capacity of gallery: {}".format(cpu_index.ntotal)) return cpu_index elif mode == 1: ngpus = faiss.get_num_gpus() print("[INFO] number of GPUs:", ngpus) res = faiss.StandardGpuResources() # use a single GPU gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) gpu_index.add(nd_feats_array) # add vectors to the index print("[INFO] Capacity of gallery: {}".format(gpu_index.ntotal)) return gpu_index elif mode == 2: multi_gpu_index = faiss.index_cpu_to_all_gpus( cpu_index) # build the index on multi GPUs multi_gpu_index.add(nd_feats_array) # add vectors to the index print("[INFO] Capacity of gallery: {}".format(multi_gpu_index.ntotal)) return multi_gpu_index
def create_index(cls, embeddings): USE_SUBSET = None if USE_SUBSET is not None: print(f">> CAREFUL. Using subset {USE_SUBSET} / {len(embeddings)}," f" {USE_SUBSET/len(embeddings):0.2%}") embeddings = embeddings[:USE_SUBSET] DIM = embeddings.shape[1] index = faiss.index_factory(DIM, FLAGS.faiss_index_factory, faiss.METRIC_INNER_PRODUCT) # index = faiss.IndexHNSWFlat(768, 128, faiss.METRIC_INNER_PRODUCT) cls.index_init(index) if FLAGS.faiss_use_gpu: print("\t- Moving to gpu") faiss_res = faiss.StandardGpuResources() index = faiss.index_cpu_to_gpu(faiss_res, 0, index) print("\t- Training the index") start = time.time() index.train(embeddings) print(f"\t- Training took " f"{tqdm.tqdm.format_interval(time.time() - start)}") print("\t- Adding the embeddings...") start = time.time() index.add(embeddings) print(f"\t- Adding took " f"{tqdm.tqdm.format_interval(time.time() - start)}") return index
def __call__(self, target_labels, features): labels, freqs = np.unique(target_labels, return_counts=True) R = len(features) faiss_search_index = faiss.IndexFlatL2(features.shape[-1]) if isinstance(features, torch.Tensor): features = features.detach().cpu().numpy() res = faiss.StandardGpuResources() faiss_search_index = faiss.index_cpu_to_gpu( res, 0, faiss_search_index) faiss_search_index.add(features) nearest_neighbours = faiss_search_index.search(features, int(R + 1))[1][:, 1:] target_labels = target_labels.reshape(-1) nn_labels = target_labels[nearest_neighbours] avg_r_precisions = [] for label, freq in zip(labels, freqs): rows_with_label = np.where(target_labels == label)[0] for row in rows_with_label: n_recalled_samples = np.arange(1, R + 1) target_label_occ_in_row = nn_labels[row, :] == label cumsum_target_label_freq_row = np.cumsum( target_label_occ_in_row) avg_r_pr_row = np.sum( cumsum_target_label_freq_row * target_label_occ_in_row / n_recalled_samples) / freq avg_r_precisions.append(avg_r_pr_row) return np.mean(avg_r_precisions)
def __init__(self, d=64, GPU=False, GPU_Number=0): #default dimension=64 self.idx = faiss.IndexFlatL2( d ) # build the index self.GPU = GPU if self.GPU: self.res = faiss.StandardGpuResources() # use a single GPU gpu = faiss.index_cpu_to_gpu(self.res, GPU_Number, self.idx) self.idx = gpu
def Indexes_of_inliers(self, Keypoints, Descriptors, buffersize): res = faiss.StandardGpuResources() nlist = 100 quantizer = faiss.IndexFlatL2(256) index = faiss.IndexIVFFlat(quantizer, 256, nlist) gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index) gpu_index_flat.train( clustering.preprocess_features(Descriptors[:buffersize])) gpu_index_flat.add( clustering.preprocess_features(Descriptors[:buffersize])) #we process the descriptors in batches of 10000 vectors rg = np.linspace(0, len(Descriptors), math.ceil(len(Descriptors) / 10000) + 1, dtype=int) keypoints_outlier_score = np.zeros(len(Keypoints)) for i in range(len(rg) - 1): descr = clustering.preprocess_features(Descriptors[rg[i]:rg[i + 1], :]) distance_to_closest_points, _ = gpu_index_flat.search(descr, 100) outlierscore = np.median(distance_to_closest_points, axis=1) keypoints_outlier_score[rg[i]:rg[i + 1]] = outlierscore inliers = keypoints_outlier_score.copy() inliers = np.sort(inliers) threshold = inliers[int( (1 - self.remove_superpoint_outliers_percentage) * (len(inliers) - 1))] inliers = keypoints_outlier_score < threshold return inliers
def _faiss_index_to_device( index: "faiss.Index", device: Optional[Union[int, List[int]]] = None) -> "faiss.Index": """ Sends a faiss index to a device. A device can either be a positive integer (GPU id), a negative integer (all GPUs), or a list of positive integers (select GPUs to use), or `None` for CPU. """ # If device is not specified, then it runs on CPU. if device is None: return index import faiss # noqa: F811 # If the device id is given as an integer if isinstance(device, int): # Positive integers are directly mapped to GPU ids if device > -1: faiss_res = faiss.StandardGpuResources() index = faiss.index_cpu_to_gpu(faiss_res, device, index) # And negative integers mean using all GPUs else: index = faiss.index_cpu_to_all_gpus(index) # Device ids given as a list mean mapping to those devices specified. elif isinstance(device, (list, tuple)): index = faiss.index_cpu_to_gpus_list(index, gpus=list(device)) else: raise TypeError( f"The argument type: {type(device)} is not expected. " + "Please pass in either nothing, a positive int, a negative int, or a list of positive ints." ) return index
def _knn_faiss(data_numpy, k, metric='euclidean', use_gpu=False): import faiss data_numpy = data_numpy.astype(np.float32) data_numpy = data_numpy.copy(order='C') data_numpy = np.ascontiguousarray(data_numpy, dtype=np.float32) if use_gpu: print('Using GPU for Faiss...') res = faiss.StandardGpuResources() else: print('Using CPU for Faiss...') if metric == 'euclidean': index = faiss.IndexFlatL2(data_numpy.shape[1]) elif metric == 'manhattan': index = faiss.IndexFlat(data_numpy.shape[1], faiss.METRIC_L1) elif metric == 'cosine': index = faiss.IndexFlat(data_numpy.shape[1], faiss.METRIC_INNER_PRODUCT) faiss.normalize_L2(data_numpy) if use_gpu: index = faiss.index_cpu_to_gpu(res, 0, index) data_numpy = np.ascontiguousarray(data_numpy, dtype=np.float32) index.train(data_numpy) assert index.is_trained index.add(data_numpy) nprobe = data_numpy.shape[0] index.nprobe = nprobe distances, neighbors = index.search(data_numpy, k) return distances, neighbors
def train_index(data, quantizer_path, trained_index_path, fine_quant='SQ8', cuda=False): quantizer = faiss.read_index(quantizer_path) if fine_quant == 'SQ8': trained_index = faiss.IndexIVFScalarQuantizer(quantizer, quantizer.d, quantizer.ntotal, faiss.METRIC_L2) elif fine_quant.startswith('PQ'): m = int(fine_quant[2:]) trained_index = faiss.IndexIVFPQ(quantizer, quantizer.d, quantizer.ntotal, m, 8) else: raise ValueError(fine_quant) if cuda: if fine_quant.startswith('PQ'): print('PQ not supported on GPU; keeping CPU.') else: res = faiss.StandardGpuResources() gpu_index = faiss.index_cpu_to_gpu(res, 0, trained_index) gpu_index.train(data) trained_index = faiss.index_gpu_to_cpu(gpu_index) else: trained_index.train(data) faiss.write_index(trained_index, trained_index_path)
def __init__(self, database, method): super().__init__(database, method) self.index = {'cosine': faiss.IndexFlatIP, 'euclidean': faiss.IndexFlatL2}[method](self.D) if os.environ.get('CUDA_VISIBLE_DEVICES'): res = faiss.StandardGpuResources() self.index = faiss.index_cpu_to_gpu(res, 0, self.index) self.add()
def get_nn_avg_dist(emb, query, knn): # cpu mode res = faiss.StandardGpuResources() # use a single GPU index = faiss.IndexFlatIP(emb.shape[1]) index = faiss.index_cpu_to_gpu(res, 0, index) index.add(emb) distances, _ = index.search(query, knn) return distances.mean(1)
def main(): parser = get_parser() args = parser.parse_args() spec = osp.basename(args.path) try: faiss_spec = parse_faiss_specs(spec.rstrip("/"))[0] except: print(spec) raise print("Faiss Spec:", faiss_spec, file=sys.stderr) if faiss_spec.pca: A = torch.from_numpy(np.load(osp.join(args.path, "pca_A.npy"))).cuda() b = torch.from_numpy(np.load(osp.join(args.path, "pca_b.npy"))).cuda() print("Loaded PCA", file=sys.stderr) centroids = np.load(osp.join(args.path, "centroids.npy")) print("Loaded centroids", centroids.shape, file=sys.stderr) res = faiss.StandardGpuResources() index_flat = (faiss.IndexFlatL2(centroids.shape[1]) if not faiss_spec.sphere else faiss.IndexFlatIP( centroids.shape[1])) faiss_index = faiss.index_cpu_to_gpu(res, 0, index_flat) faiss_index.add(centroids) generator, num, root = get_iterator(args) iterator = generator() had_labels = False label_path = osp.join(args.path, f"{args.split}.{args.labels}") with torch.no_grad(): with open(osp.join(args.path, f"{args.split}.src"), "w") as fp, open(osp.join(args.path, f"{args.split}.tsv"), "w") as pp, open(label_path, "w") as lp: print(root, file=pp) for f, fname, lbl in tqdm.tqdm(iterator, total=num): if faiss_spec.pca: f = torch.mm(f, A) + b if faiss_spec.norm: f = F.normalize(f, p=2, dim=-1) f = f.cpu().numpy() _, z = faiss_index.search(f, 1) print(" ".join(str(x.item()) for x in z), file=fp) print(fname, file=pp) if lbl is not None: print(lbl, file=lp) had_labels = True if not had_labels: os.remove(label_path)
def add_vectors( self, vectors: Union[np.array, "Dataset"], column: Optional[str] = None, batch_size: int = 1000, train_size: Optional[int] = None, faiss_verbose: Optional[bool] = None, ): """ Add vectors to the index. If the arrays are inside a certain column, you can specify it using the `column` argument. """ import faiss # noqa: F811 # Create index if self.faiss_index is None: size = len(vectors[0]) if column is None else len(vectors[0][column]) if self.string_factory is not None: if self.metric_type is None: index = faiss.index_factory(size, self.string_factory) else: index = faiss.index_factory(size, self.string_factory, self.metric_type) else: if self.metric_type is None: index = faiss.IndexFlat(size) else: index = faiss.IndexFlat(size, self.metric_type) if self.device is not None and self.device > -1: self.faiss_res = faiss.StandardGpuResources() index = faiss.index_cpu_to_gpu(self.faiss_res, self.device, index) self.faiss_index = index logger.info(f"Created faiss index of type {type(self.faiss_index)}") # Set verbosity level if faiss_verbose is not None: self.faiss_index.verbose = faiss_verbose if hasattr(self.faiss_index, "index") and self.faiss_index.index is not None: self.faiss_index.index.verbose = faiss_verbose if hasattr(self.faiss_index, "quantizer") and self.faiss_index.quantizer is not None: self.faiss_index.quantizer.verbose = faiss_verbose if hasattr(self.faiss_index, "clustering_index") and self.faiss_index.clustering_index is not None: self.faiss_index.clustering_index.verbose = faiss_verbose # Train if train_size is not None: train_vecs = vectors[:train_size] if column is None else vectors[:train_size][column] logger.info(f"Training the index with the first {len(train_vecs)} vectors") self.faiss_index.train(train_vecs) else: logger.info("Ignored the training step of the faiss index as `train_size` is None.") # Add vectors logger.info(f"Adding {len(vectors)} vectors to the faiss index") for i in utils.tqdm( range(0, len(vectors), batch_size), disable=bool(logging.get_verbosity() == logging.NOTSET) ): vecs = vectors[i : i + batch_size] if column is None else vectors[i : i + batch_size][column] self.faiss_index.add(vecs)
def indexer(self): index = faiss.index_factory(self.dimensions, INDEX_KEY) if USE_GPU: res = faiss.StandardGpuResources() index = faiss.index_cpu_to_gpu(res, 0, index) images_list = self.iterate_files() # prepare ids ids_count = 0 index_dict = {} ids = None features = np.matrix([]) for file_name in images_list: ret, sift_feature = self.calc_sift(file_name) if ret == 0 and sift_feature.any(): # record id and path image_dict = {ids_count: (file_name, sift_feature)} index_dict.update(image_dict) ids_list = np.linspace(ids_count, ids_count, num=sift_feature.shape[0], dtype="int64") ids_count += 1 if features.any(): features = np.vstack((features, sift_feature)) ids = np.hstack((ids, ids_list)) else: features = sift_feature ids = ids_list if ids_count % 500 == 499: if not index.is_trained and INDEX_KEY != "IDMap,Flat": index.train(features) index.add_with_ids(features, ids) ids = None features = np.matrix([]) if features.any(): print("training..") if not index.is_trained and INDEX_KEY != "IDMap,Flat": index.train(features) index.add_with_ids(features, ids) # save index print("saving index..") faiss.write_index(index, INDEX_PATH) # save ids with open(IDS_VECTORS_PATH, "wb+") as f: try: pickle.dump(index_dict, f, True) except EnvironmentError as e: print("Failed to save index file error:[{}]".format(e)) f.close() except RuntimeError as v: print("Failed to save index file error:[{}]".format(v)) f.close() # print("N", index.ntotal, dir(index), index.__dict__) return index.ntotal
def __init__(self, dimensions, gpu=False): if gpu: # requires faiss-gpu res = faiss.StandardGpuResources() # use a single GPU index_flat = faiss.IndexFlatIP(dimensions) # build a flat index self.index = faiss.index_cpu_to_gpu(res, 0, index_flat) else: self.index = faiss.IndexFlatIP(dimensions) self.ids = []
def __init__(self, x, gpu_id, verbose=False): DatasetAssign.__init__(self, x) index = faiss.IndexFlatL2(x.shape[1]) if gpu_id >= 0: self.index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), gpu_id, index) else: # -1 -> assign to all GPUs self.index = faiss.index_cpu_to_all_gpus(index)
def samplePC(cubes, flip_bbox=False, split_bbox=False): cube_geom = [] for c in cubes: cube_geom.append( torch.cat( (c['xd'].unsqueeze(0), c['yd'].unsqueeze(0), c['zd'].unsqueeze(0), c['center'], c['xdir'], c['ydir']))) scene_geom = torch.stack([c for c in cube_geom]).to(device) ind_to_pc = {} for i in range(0, scene_geom.shape[0]): xyz = s_xyz s_inds = (torch.ones(1, xyz.shape[1]) * i).long().to(device) s_r = torch.cat( ((scene_geom[s_inds][:, :, 6:9] / (scene_geom[s_inds][:, :, 6:9].norm(dim=2).unsqueeze(2) + 1e-8)).unsqueeze(3), (scene_geom[s_inds][:, :, 9:12] / (scene_geom[s_inds][:, :, 9:12].norm(dim=2).unsqueeze(2) + 1e-8)).unsqueeze(3), torch.cross( scene_geom[s_inds][:, :, 6:9] / (scene_geom[s_inds][:, :, 6:9].norm(dim=2).unsqueeze(2) + 1e-8), scene_geom[s_inds][:, :, 9:12] / (scene_geom[s_inds][:, :, 9:12].norm(dim=2).unsqueeze(2) + 1e-8)).unsqueeze(3)), dim=3) s_out = ((s_r @ (( (xyz - .5) * scene_geom[s_inds][:, :, :3]).unsqueeze(-1)) ).squeeze() + scene_geom[s_inds][:, :, 3:6]).squeeze() ind_to_pc[i] = s_out if flip_bbox: ind_to_pc[0] += bb_mask temp = ind_to_pc[0].clone() ind_to_pc[0][ft] = temp[fbo] ind_to_pc[0][fbo] = temp[ft] if split_bbox: bbox_pc = ind_to_pc.pop(0) ind_to_pc[-2] = bbox_pc.clone() + bot_mask ind_to_pc[-1] = bbox_pc.clone() + top_mask res = {} for key in ind_to_pc: index_cpu = faiss.IndexFlatL2(3) index = faiss.index_cpu_to_gpu(resource, torch.cuda.current_device(), index_cpu) index.add(np.ascontiguousarray(ind_to_pc[key].cpu().numpy())) res[key] = (ind_to_pc[key], index) return res, scene_geom
def fit(self, X): X = X.astype(numpy.float32) self._index = faiss.index_factory(len(X[0]), "IVF%d,PQ64" % self._n_bits) co = faiss.GpuClonerOptions() co.useFloat16 = True self._index = faiss.index_cpu_to_gpu(self._res, 0, self._index, co) self._index.train(X) self._index.add(X) self._index.setNumProbes(self._n_probes)
def make_index(self, flags_obj): self.make_index_brute_force(flags_obj) if flags_obj.cg_use_gpu: res = faiss.StandardGpuResources() self.index = faiss.index_cpu_to_gpu(res, flags_obj.cg_gpu_id, self.index)
def do_cpu_to_gpu(self, index_key): ts = [] ts.append(time.time()) (xt, xb, xq) = self.get_dataset(small_one=True) nb, d = xb.shape index = faiss.index_factory(d, index_key) if index.__class__ == faiss.IndexIVFPQ: # speed up test index.pq.cp.niter = 2 index.do_polysemous_training = False ts.append(time.time()) index.train(xt) ts.append(time.time()) # adding some ids because there was a bug in this case index.add_with_ids(xb, np.arange(nb) * 3 + 12345) ts.append(time.time()) index.nprobe = 4 D, Iref = index.search(xq, 10) ts.append(time.time()) res = faiss.StandardGpuResources() gpu_index = faiss.index_cpu_to_gpu(res, 0, index) ts.append(time.time()) gpu_index.setNumProbes(4) D, Inew = gpu_index.search(xq, 10) ts.append(time.time()) print 'times:', [t - ts[0] for t in ts] self.assertGreaterEqual((Iref == Inew).sum(), Iref.size) if faiss.get_num_gpus() == 1: return for shard in False, True: # test on just 2 GPUs res = [faiss.StandardGpuResources() for i in range(2)] co = faiss.GpuMultipleClonerOptions() co.shard = shard gpu_index = faiss.index_cpu_to_gpu_multiple_py(res, index, co) faiss.GpuParameterSpace().set_index_parameter( gpu_index, 'nprobe', 4) D, Inew = gpu_index.search(xq, 10) self.assertGreaterEqual((Iref == Inew).sum(), Iref.size)
# keep track of optimal operating points seen so far op = faiss.OperatingPoints() for index_key in keys_to_test: print "============ key", index_key # make the index described by the key index = faiss.index_factory(d, index_key) if use_gpu: # transfer to GPU (may be partial) index = faiss.index_cpu_to_gpu(res, dev_no, index) params = faiss.GpuParameterSpace() else: params = faiss.ParameterSpace() params.initialize(index) print "[%.3f s] train & add" % (time.time() - t0) index.train(xt) index.add(xb) print "[%.3f s] explore op points" % (time.time() - t0) # find operating points for this index opi = params.explore(index, xq, crit)
def test_set_gpu_param(self): index = faiss.index_factory(12, "PCAR8,IVF10,PQ4") res = faiss.StandardGpuResources() gpu_index = faiss.index_cpu_to_gpu(res, 0, index) faiss.GpuParameterSpace().set_index_parameter(gpu_index, "nprobe", 3)
# need to contactinate tests as well xb = np.load("train_filenames.txt") xq = np.load("train_featues.txt") print(xb.shape) print(xq.shape) print(xq.shape) import faiss res = faiss.StandardGpuResources() # use a single GPU # build a flat (CPU) index index_flat = faiss.IndexFlatL2(d) # make it into a gpu index gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat) gpu_index_flat.add(xb) # add vectors to the index print(gpu_index_flat.ntotal) k = 100 # we want to see 4 nearest neighbors D, I = gpu_index_flat.search(xq, k) # actual search # print(I[:5]) # neighbors of the 5 first queries # print(I[-5:]) # neighbors of the 5 last queries np.save("output/I.npy", I) np.save("output/D.npy", D) ### make submission index_path = "input/index/" index_list = sorted(glob.glob(index_path + "*")) # 1091756
################################################################# print "============ Approximate search" index = faiss.index_factory(d, "IVF4096,PQ64") # faster, uses more memory # index = faiss.index_factory(d, "IVF16384,Flat") co = faiss.GpuClonerOptions() # here we are using a 64-byte PQ, so we must set the lookup tables to # 16 bit float (this is due to the limited temporary memory). co.useFloat16 = True index = faiss.index_cpu_to_gpu(res, 0, index, co) print "train" index.train(xt) print "add vectors to index" index.add(xb) print "warmup" index.search(xq, 123) print "benchmark"