All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
test_pairs_decoding.cpp
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <cstdio>
#include <cstdlib>
#include <memory>
#include <vector>
#include <gtest/gtest.h>
#include <faiss/IndexIVF.h>
#include <faiss/AutoTune.h>
#include <faiss/VectorTransform.h>
/*************************************************************
* The functions to test, that can be useful in FANN
*************************************************************/
/* Returns the cluster the embeddings belong to.
*
* @param index Index, which should be an IVF index
* (otherwise there are no clusters)
* @param embeddings object descriptors for which the centroids should be found,
* size num_objects * d
* @param cebtroid_ids
* cluster id each object belongs to, size num_objects
*/
void Search_centroid(faiss::Index *index,
const float* embeddings, int num_objects,
int64_t* centroid_ids)
{
const float *x = embeddings;
std::unique_ptr<float> del;
if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
x = index_pre->apply_chain(num_objects, x);
del.reset((float*)x);
index = index_pre->index;
}
faiss::IndexIVF* index_ivf = dynamic_cast< faiss::IndexIVF*>(index);
assert(index_ivf);
index_ivf->quantizer->assign(num_objects, x, centroid_ids);
}
/* Returns the cluster the embeddings belong to.
*
* @param index Index, which should be an IVF index
* (otherwise there are no clusters)
* @param query_centroid_ids
* centroid ids corresponding to the query vectors (size n)
* @param result_centroid_ids
* centroid ids corresponding to the results (size n * k)
* other arguments are the same as the standard search function
*/
void search_and_retrun_centroids(faiss::Index *index,
size_t n,
const float* xin,
long k,
float *distances,
int64_t* labels,
int64_t* query_centroid_ids,
int64_t* result_centroid_ids)
{
const float *x = xin;
std::unique_ptr<float> del;
if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
x = index_pre->apply_chain(n, x);
del.reset((float*)x);
index = index_pre->index;
}
faiss::IndexIVF* index_ivf = dynamic_cast< faiss::IndexIVF*>(index);
assert(index_ivf);
size_t nprobe = index_ivf->nprobe;
std::vector<long> cent_nos (n * nprobe);
std::vector<float> cent_dis (n * nprobe);
index_ivf->quantizer->search(
n, x, nprobe, cent_dis.data(), cent_nos.data());
if (query_centroid_ids) {
for (size_t i = 0; i < n; i++)
query_centroid_ids[i] = cent_nos[i * nprobe];
}
index_ivf->search_preassigned (n, x, k,
cent_nos.data(), cent_dis.data(),
distances, labels, true);
for (size_t i = 0; i < n * k; i++) {
int64_t label = labels[i];
if (label < 0) {
if (result_centroid_ids)
result_centroid_ids[i] = -1;
} else {
long list_no = label >> 32;
long list_index = label & 0xffffffff;
if (result_centroid_ids)
result_centroid_ids[i] = list_no;
labels[i] = index_ivf->ids[list_no][list_index];
}
}
}
/*************************************************************
* Test utils
*************************************************************/
// return an IndexIVF that may be embedded in an IndexPreTransform
faiss::IndexIVF * get_IndexIVF(faiss::Index *index) {
if (auto index_pre = dynamic_cast<faiss::IndexPreTransform*>(index)) {
index = index_pre->index;
}
faiss::IndexIVF* index_ivf = dynamic_cast< faiss::IndexIVF*>(index);
bool t = index_ivf != nullptr;
assert(index_ivf);
return index_ivf;
}
// dimension of the vectors to index
int d = 64;
// size of the database we plan to index
size_t nb = 8000;
// nb of queries
size_t nq = 200;
std::vector<float> make_data(size_t n)
{
std::vector <float> database (n * d);
for (size_t i = 0; i < n * d; i++) {
database[i] = drand48();
}
return database;
}
std::unique_ptr<faiss::Index> make_index(const char *index_type,
const std::vector<float> & x) {
auto index = std::unique_ptr<faiss::Index> (
faiss::index_factory(d, index_type));
index->train(nb, x.data());
index->add(nb, x.data());
return index;
}
/*************************************************************
* Test functions for a given index type
*************************************************************/
bool test_Search_centroid(const char *index_key) {
std::vector<float> xb = make_data(nb); // database vectors
auto index = make_index(index_key, xb);
/* First test: find the centroids associated to the database
vectors and make sure that each vector does indeed appear in
the inverted list corresponding to its centroid */
std::vector<int64_t> centroid_ids (nb);
Search_centroid(index.get(), xb.data(), nb, centroid_ids.data());
const faiss::IndexIVF * ivf = get_IndexIVF(index.get());
for(int i = 0; i < nb; i++) {
bool found = false;
int list_no = centroid_ids[i];
for(int j: ivf->ids[list_no]) {
if (j == i) {
found = true;
break;
}
}
if(!found) return false;
}
return true;
}
int test_search_and_return_centroids(const char *index_key) {
std::vector<float> xb = make_data(nb); // database vectors
auto index = make_index(index_key, xb);
std::vector<int64_t> centroid_ids (nb);
Search_centroid(index.get(), xb.data(), nb, centroid_ids.data());
faiss::IndexIVF * ivf = get_IndexIVF(index.get());
ivf->nprobe = 4;
std::vector<float> xq = make_data(nq); // database vectors
int k = 5;
// compute a reference search result
std::vector<long> refI (nq * k);
std::vector<float> refD (nq * k);
index->search (nq, xq.data(), k, refD.data(), refI.data());
// compute search result
std::vector<long> newI (nq * k);
std::vector<float> newD (nq * k);
std::vector<int64_t> query_centroid_ids (nq);
std::vector<int64_t> result_centroid_ids (nq * k);
search_and_retrun_centroids(index.get(),
nq, xq.data(), k,
newD.data(), newI.data(),
query_centroid_ids.data(),
result_centroid_ids.data());
// first verify that we have the same result as the standard search
if (newI != refI) {
return 1;
}
// then check if the result ids are indeed in the inverted list
// they are supposed to be in
for(int i = 0; i < nq * k; i++) {
int list_no = result_centroid_ids[i];
int result_no = newI[i];
if (result_no < 0) continue;
bool found = false;
for(int j: ivf->ids[list_no]) {
if (j == result_no) {
found = true;
break;
}
}
if(!found) return 2;
}
return 0;
}
/*************************************************************
* Test entry points
*************************************************************/
TEST(test_Search_centroid, IVFFlat) {
bool ok = test_Search_centroid("IVF32,Flat");
EXPECT_TRUE(ok);
}
TEST(test_Search_centroid, PCAIVFFlat) {
bool ok = test_Search_centroid("PCA16,IVF32,Flat");
EXPECT_TRUE(ok);
}
TEST(test_search_and_return_centroids, IVFFlat) {
int err = test_search_and_return_centroids("IVF32,Flat");
EXPECT_NE(err, 1);
EXPECT_NE(err, 2);
}
TEST(test_search_and_return_centroids, PCAIVFFlat) {
int err = test_search_and_return_centroids("PCA16,IVF32,Flat");
EXPECT_NE(err, 1);
EXPECT_NE(err, 2);
}
faiss::IndexIVF::search_preassigned
virtual void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs) const =0
faiss::IndexIVF
Definition: IndexIVF.h:45
faiss::IndexIVF::nprobe
size_t nprobe
number of probes at query time
Definition: IndexIVF.h:47
faiss::Index::train
virtual void train(idx_t, const float *)
Definition: Index.h:89
faiss::Index::assign
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
Definition: Index.cpp:23
faiss::IndexIVF::ids
std::vector< std::vector< long > > ids
Inverted lists for indexes.
Definition: IndexIVF.h:62
faiss::IndexIVF::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:49
faiss::Index::add
virtual void add(idx_t n, const float *x)=0
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::Index
Definition: Index.h:60
faiss::index_factory
Index * index_factory(int d, const char *description_in, MetricType metric)
Definition: AutoTune.cpp:668