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)
