def do_test(self, metric, metric_arg=0): res = faiss.StandardGpuResources() d = 32 nb = 1000 nq = 100 rs = np.random.RandomState(123) xb = rs.rand(nb, d).astype('float32') xq = rs.rand(nq, d).astype('float32') index_ref = faiss.IndexFlat(d, metric) index_ref.metric_arg = metric_arg index_ref.add(xb) Dref, Iref = index_ref.search(xq, 10) # build from other index index = faiss.GpuIndexFlat(res, index_ref) Dnew, Inew = index.search(xq, 10) np.testing.assert_array_equal(Inew, Iref) np.testing.assert_allclose(Dnew, Dref, rtol=1e-6) # build from scratch index = faiss.GpuIndexFlat(res, d, metric) index.metric_arg = metric_arg index.add(xb) Dnew, Inew = index.search(xq, 10) np.testing.assert_array_equal(Inew, Iref)
def _set_block_index(self): """Create a Faiss Flat index with inner product as the metric to search against""" try: import faiss except ImportError: raise Exception( "Error: Please install faiss to use FaissMIPSIndex") if mpu.is_unitialized() or mpu.get_data_parallel_rank() == 0: print("\n> Building index", flush=True) self.block_mips_index = faiss.index_factory(self.embed_size, 'Flat', faiss.METRIC_INNER_PRODUCT) if self.use_gpu: # create resources and config for GpuIndex res = faiss.StandardGpuResources() config = faiss.GpuIndexFlatConfig() config.device = torch.cuda.current_device() config.useFloat16 = True self.block_mips_index = faiss.GpuIndexFlat(res, self.block_mips_index, config) if mpu.is_unitialized() or mpu.get_data_parallel_rank() == 0: print(">> Initialized index on GPU {}".format( self.block_mips_index.getDevice()), flush=True) else: # CPU index supports IDs so wrap with IDMap self.block_mips_index = faiss.IndexIDMap(self.block_mips_index) if mpu.is_unitialized() or mpu.get_data_parallel_rank() == 0: print(">> Initialized index on CPU", flush=True) # if we were constructed with a BlockData, then automatically load it when the FAISS structure is built if self.block_data is not None: self.add_block_embed_data(self.block_data)
def test_flat(self): index = faiss.GpuIndexFlat(faiss.StandardGpuResources(), self.d, faiss.METRIC_L2) index.add(self.xb)
def fit(self, X): X = X.astype(numpy.float32) self._index = faiss.GpuIndexFlat(self._res, len(X[0]), faiss.METRIC_L2) self._index.train(X) self._index.add(X)