def test_indices_ivfpq(self):
res = faiss.StandardGpuResources()
d = 128
nb = 5000
nlist = 10
M = 4
nbits = 8
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xb_indices_base = np.arange(nb, dtype=np.int64)
# Force values to not be representable in int32
xb_indices = (xb_indices_base + 4294967296).astype('int64')
config = faiss.GpuIndexIVFPQConfig()
idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits, faiss.METRIC_L2,
config)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
_, I = idx.search(xb[10:20], 5)
self.assertTrue(np.array_equal(xb_indices[10:20], I[:, 0]))
# Store values using 32-bit indices instead
config.indicesOptions = faiss.INDICES_32_BIT
idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits, faiss.METRIC_L2,
config)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
_, I = idx.search(xb[10:20], 5)
# This will strip the high bit
self.assertTrue(np.array_equal(xb_indices_base[10:20], I[:, 0]))
def test_IndexIVFPQ(self):
(xt, xb, xq) = self.get_dataset()
d = xt.shape[1]
dev_no = 0
usePrecomputed = True
res = faiss.StandardGpuResources()
gt_index = faiss.GpuIndexFlatL2(res, dev_no, d, False)
gt_index.add(xb)
D, gt_nns = gt_index.search(xq, 1)
coarse_quantizer = faiss.IndexFlatL2(d)
ncentroids = int(np.sqrt(xb.shape[0])) * 4
index = faiss.IndexIVFPQ(coarse_quantizer, d, ncentroids, 32, 8)
# add implemented on GPU but not train
index.train(xt)
gpuIndex = faiss.GpuIndexIVFPQ(res, dev_no, faiss.INDICES_64_BIT,
False, index)
gpuIndex.setPrecomputedCodes(usePrecomputed)
gpuIndex.setNumProbes(64)
index.add(xb)
D, nns = index.search(xq, 10)
n_ok = (nns == gt_nns).sum()
nq = xq.shape[0]
print ncentroids, n_ok, nq
self.assertGreater(n_ok, nq * 0.2)
def test_ivfpq(self):
index_cpu = faiss.IndexIVFPQ(faiss.IndexFlatL2(self.d), self.d,
self.nlist, 2, 8)
# speed up test
index_cpu.pq.cp.niter = 2
index_cpu.do_polysemous_training = False
index_cpu.train(self.xb)
index = faiss.GpuIndexIVFPQ(faiss.StandardGpuResources(), index_cpu)
index.add(self.xb)
def test_serialize(self):
res = faiss.StandardGpuResources()
d = 32
k = 10
train = make_t(10000, d)
add = make_t(10000, d)
query = make_t(10, d)
# Construct various GPU index types
indexes = []
# Flat
indexes.append(faiss.GpuIndexFlatL2(res, d))
# IVF
nlist = 5
# IVFFlat
indexes.append(faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2))
# IVFSQ
indexes.append(faiss.GpuIndexIVFScalarQuantizer(res, d, nlist, faiss.ScalarQuantizer.QT_fp16))
# IVFPQ
indexes.append(faiss.GpuIndexIVFPQ(res, d, nlist, 4, 8, faiss.METRIC_L2))
for index in indexes:
index.train(train)
index.add(add)
orig_d, orig_i = index.search(query, k)
ser = faiss.serialize_index(faiss.index_gpu_to_cpu(index))
cpu_index = faiss.deserialize_index(ser)
gpu_index_restore = faiss.index_cpu_to_gpu(res, 0, cpu_index)
restore_d, restore_i = gpu_index_restore.search(query, k)
self.assertTrue(np.array_equal(orig_d, restore_d))
self.assertTrue(np.array_equal(orig_i, restore_i))
# Make sure the index is in a state where we can add to it
# without error
gpu_index_restore.add(query)
def test_copy_to_gpu(self):
res = faiss.StandardGpuResources()
for bits_per_code in [4, 5, 6, 8]:
d = 128
nb = 10000
nq = 20
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
nlist = int(math.sqrt(nb))
sub_q = 16
bits_per_code = 8
nprobe = 4
config = faiss.GpuIndexIVFPQConfig()
config.interleavedLayout = True
idx_gpu = faiss.GpuIndexIVFPQ(res, d, nlist, sub_q, bits_per_code,
faiss.METRIC_L2, config)
q = faiss.IndexFlatL2(d)
idx_cpu = faiss.IndexIVFPQ(q, d, nlist, sub_q, bits_per_code,
faiss.METRIC_L2)
idx_cpu.train(xb)
idx_cpu.add(xb)
idx_gpu.copyFrom(idx_cpu)
idx_gpu.nprobe = nprobe
idx_cpu.nprobe = nprobe
# Try without precomputed codes
d_g, i_g = idx_gpu.search(xq, 10)
d_c, i_c = idx_cpu.search(xq, 10)
self.assertGreaterEqual((i_g == i_c).sum(), i_g.size * 0.9)
self.assertTrue(np.allclose(d_g, d_c))
# Try with precomputed codes (different kernel)
idx_gpu.setPrecomputedCodes(True)
d_g, i_g = idx_gpu.search(xq, 10)
d_c, i_c = idx_cpu.search(xq, 10)
self.assertGreaterEqual((i_g == i_c).sum(), i_g.size * 0.9)
self.assertTrue(np.allclose(d_g, d_c))
def test_indices_ivfpq(self):
res = faiss.StandardGpuResources()
d = 128
nb = 5000
nlist = 10
M = 4
nbits = 8
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xb_indices_base = np.arange(nb, dtype=np.int64)
# Force values to not be representable in int32
xb_indices = (xb_indices_base + 4294967296).astype('int64')
config = faiss.GpuIndexIVFPQConfig()
idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits,
faiss.METRIC_L2, config)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
# invalid k (should be > 0)
k = -5
idx.setNumProbes(3)
self.assertRaises(AssertionError, idx.search, xb[10:20], k)
# invalid nprobe (should be > 0)
self.assertRaises(RuntimeError, idx.setNumProbes, 0)
self.assertRaises(RuntimeError, idx.setNumProbes, -3)
k = 5
idx.nprobe = -3
self.assertRaises(RuntimeError, idx.search, xb[10:20], k)
# valid params
k = 5
idx.setNumProbes(3)
_, I = idx.search(xb[10:20], k)
self.assertTrue(np.array_equal(xb_indices[10:20], I[:, 0]))
def test_IndexIVFPQ(self):
(xt, xb, xq) = self.get_dataset()
d = xt.shape[1]
dev_no = 0
usePrecomputed = True
res = faiss.StandardGpuResources()
flat_config = faiss.GpuIndexFlatConfig()
flat_config.device = dev_no
gt_index = faiss.GpuIndexFlatL2(res, d, flat_config)
gt_index.add(xb)
D, gt_nns = gt_index.search(xq, 1)
coarse_quantizer = faiss.IndexFlatL2(d)
ncentroids = int(np.sqrt(xb.shape[0])) * 4
index = faiss.IndexIVFPQ(coarse_quantizer, d, ncentroids, 32, 8)
# add implemented on GPU but not train
index.train(xt)
ivfpq_config = faiss.GpuIndexIVFPQConfig()
ivfpq_config.device = dev_no
ivfpq_config.usePrecomputedTables = usePrecomputed
gpuIndex = faiss.GpuIndexIVFPQ(res, index, ivfpq_config)
gpuIndex.setNumProbes(64)
index.add(xb)
D, nns = index.search(xq, 10)
n_ok = (nns == gt_nns).sum()
nq = xq.shape[0]
print ncentroids, n_ok, nq
self.assertGreater(n_ok, nq * 0.2)
start = time.time()
mfcc_features = np.array(mfcc_features, dtype=np.float32)
labels = np.array(labels)
print("formatting labels took : ", time.time() - start, " seconds")
### Build the FAISS index - in this case, IVFPQ
d = 16
print('starting training')
start = time.time()
res = faiss.StandardGpuResources() # use a single GPU
ncentroids = int(4 * math.sqrt(
len(labels))) #ncentroids defined on the low end of FAISS guidance.
gpu_index = faiss.GpuIndexIVFPQ(res, 16, ncentroids, 8, 8, faiss.METRIC_L2)
gpu_index.train(mfcc_features)
gpu_index.add_with_ids(
mfcc_features, labels
) # adding vectors to the index w/ ids means that the index handles label lookups
print("building index took : ", time.time() - start, " seconds")
#save the index to file
index = faiss.index_gpu_to_cpu(gpu_index)
faiss.write_index(index, "/home/sir/voice/faiss_index_librispeech360.ind")
### Read the test files into GPU memory (and do not bring them to the CPU) to take advantage of GPU querying of torch tensors
start = time.time()
calc_mfcc = torchaudio.transforms.MFCC(n_mfcc=13,
