All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
test_params_override.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/IndexBinaryIVF.h>
#include <faiss/AutoTune.h>
#include <faiss/IVFlib.h>
using namespace faiss;
namespace {
typedef Index::idx_t idx_t;
// dimension of the vectors to index
int d = 32;
// size of the database we plan to index
size_t nb = 1000;
// 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<Index> make_index(const char *index_type,
MetricType metric,
const std::vector<float> & x)
{
std::unique_ptr<Index> index(index_factory(d, index_type, metric));
index->train(nb, x.data());
index->add(nb, x.data());
return index;
}
std::vector<idx_t> search_index(Index *index, const float *xq) {
int k = 10;
std::vector<idx_t> I(k * nq);
std::vector<float> D(k * nq);
index->search (nq, xq, k, D.data(), I.data());
return I;
}
std::vector<idx_t> search_index_with_params(
Index *index, const float *xq, IVFSearchParameters *params) {
int k = 10;
std::vector<idx_t> I(k * nq);
std::vector<float> D(k * nq);
ivflib::search_with_parameters (index, nq, xq, k,
D.data(), I.data(), params);
return I;
}
/*************************************************************
* Test functions for a given index type
*************************************************************/
int test_params_override (const char *index_key, MetricType metric) {
std::vector<float> xb = make_data(nb); // database vectors
auto index = make_index(index_key, metric, xb);
//index->train(nb, xb.data());
// index->add(nb, xb.data());
std::vector<float> xq = make_data(nq);
ParameterSpace ps;
ps.set_index_parameter(index.get(), "nprobe", 2);
auto res2ref = search_index(index.get(), xq.data());
ps.set_index_parameter(index.get(), "nprobe", 9);
auto res9ref = search_index(index.get(), xq.data());
ps.set_index_parameter(index.get(), "nprobe", 1);
IVFSearchParameters params;
params.max_codes = 0;
params.nprobe = 2;
auto res2new = search_index_with_params(index.get(), xq.data(), &params);
params.nprobe = 9;
auto res9new = search_index_with_params(index.get(), xq.data(), &params);
if (res2ref != res2new)
return 2;
if (res9ref != res9new)
return 9;
return 0;
}
} // namespace
/*************************************************************
* Test entry points
*************************************************************/
TEST(TPO, IVFFlat) {
int err1 = test_params_override ("IVF32,Flat", METRIC_L2);
EXPECT_EQ(err1, 0);
int err2 = test_params_override ("IVF32,Flat", METRIC_INNER_PRODUCT);
EXPECT_EQ(err2, 0);
}
TEST(TPO, IVFPQ) {
int err1 = test_params_override ("IVF32,PQ8np", METRIC_L2);
EXPECT_EQ(err1, 0);
int err2 = test_params_override ("IVF32,PQ8np", METRIC_INNER_PRODUCT);
EXPECT_EQ(err2, 0);
}
TEST(TPO, IVFSQ) {
int err1 = test_params_override ("IVF32,SQ8", METRIC_L2);
EXPECT_EQ(err1, 0);
int err2 = test_params_override ("IVF32,SQ8", METRIC_INNER_PRODUCT);
EXPECT_EQ(err2, 0);
}
TEST(TPO, IVFFlatPP) {
int err1 = test_params_override ("PCA16,IVF32,SQ8", METRIC_L2);
EXPECT_EQ(err1, 0);
int err2 = test_params_override ("PCA16,IVF32,SQ8", METRIC_INNER_PRODUCT);
EXPECT_EQ(err2, 0);
}
/*************************************************************
* Same for binary indexes
*************************************************************/
std::vector<uint8_t> make_data_binary(size_t n) {
std::vector <uint8_t> database (n * d / 8);
for (size_t i = 0; i < n * d / 8; i++) {
database[i] = lrand48();
}
return database;
}
std::unique_ptr<IndexBinaryIVF> make_index(const char *index_type,
const std::vector<uint8_t> & x)
{
auto index = std::unique_ptr<IndexBinaryIVF>
(dynamic_cast< IndexBinaryIVF*>(index_binary_factory (d, index_type)));
index->train(nb, x.data());
index->add(nb, x.data());
return index;
}
std::vector<idx_t> search_index(IndexBinaryIVF *index, const uint8_t *xq) {
int k = 10;
std::vector<idx_t> I(k * nq);
std::vector<int32_t> D(k * nq);
index->search (nq, xq, k, D.data(), I.data());
return I;
}
std::vector<idx_t> search_index_with_params(
IndexBinaryIVF *index, const uint8_t *xq, IVFSearchParameters *params) {
int k = 10;
std::vector<idx_t> I(k * nq);
std::vector<int32_t> D(k * nq);
std::vector<idx_t> Iq(params->nprobe * nq);
std::vector<int32_t> Dq(params->nprobe * nq);
index->quantizer->search(nq, xq, params->nprobe,
Dq.data(), Iq.data());
index->search_preassigned(nq, xq, k, Iq.data(), Dq.data(),
D.data(), I.data(),
false, params);
return I;
}
int test_params_override_binary (const char *index_key) {
std::vector<uint8_t> xb = make_data_binary(nb); // database vectors
auto index = make_index (index_key, xb);
index->train(nb, xb.data());
index->add(nb, xb.data());
std::vector<uint8_t> xq = make_data_binary(nq);
index->nprobe = 2;
auto res2ref = search_index(index.get(), xq.data());
index->nprobe = 9;
auto res9ref = search_index(index.get(), xq.data());
index->nprobe = 1;
IVFSearchParameters params;
params.max_codes = 0;
params.nprobe = 2;
auto res2new = search_index_with_params(index.get(), xq.data(), &params);
params.nprobe = 9;
auto res9new = search_index_with_params(index.get(), xq.data(), &params);
if (res2ref != res2new)
return 2;
if (res9ref != res9new)
return 9;
return 0;
}
TEST(TPOB, IVF) {
int err1 = test_params_override_binary ("BIVF32");
EXPECT_EQ(err1, 0);
}
faiss::IndexShards::train
void train(idx_t n, const float *x) override
Definition: MetaIndexes.cpp:399
faiss::IndexBinary::search
virtual void search(idx_t n, const uint8_t *x, idx_t k, int32_t *distances, idx_t *labels) const =0
faiss::IndexBinaryIVF::nprobe
size_t nprobe
number of probes at query time
Definition: IndexBinaryIVF.h:41
faiss::IndexBinaryIVF
Definition: IndexBinaryIVF.h:36
faiss::IndexBinaryIVF::search_preassigned
void search_preassigned(idx_t n, const uint8_t *x, idx_t k, const idx_t *assign, const int32_t *centroid_dis, int32_t *distances, idx_t *labels, bool store_pairs, const IVFSearchParameters *params=nullptr) const
Definition: IndexBinaryIVF.cpp:662
faiss::Index::train
virtual void train(idx_t n, const float *x)
Definition: Index.cpp:24
faiss::IndexBinaryIVF::quantizer
IndexBinary * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexBinaryIVF.h:53
faiss::IVFSearchParameters
Definition: IndexIVF.h:61
faiss::IndexBinaryIVF::train
void train(idx_t n, const uint8_t *x) override
Trains the quantizer and calls train_residual to train sub-quantizers.
Definition: IndexBinaryIVF.cpp:250
faiss::Index::add
virtual void add(idx_t n, const float *x)=0
faiss::IndexShards::add
void add(idx_t n, const float *x) override
supported only for sub-indices that implement add_with_ids
Definition: MetaIndexes.cpp:422
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::IVFSearchParameters::nprobe
size_t nprobe
number of probes at query time
Definition: IndexIVF.h:62
faiss::IndexBinaryIVF::add
void add(idx_t n, const uint8_t *x) override
Quantizes x and calls add_with_key.
Definition: IndexBinaryIVF.cpp:57
faiss::ParameterSpace::set_index_parameter
virtual void set_index_parameter(Index *index, const std::string &name, double val) const
set one of the parameters
Definition: AutoTune.cpp:452
faiss::ParameterSpace
Definition: AutoTune.h:132
faiss::IndexBinaryIVF::search
virtual void search(idx_t n, const uint8_t *x, idx_t k, int32_t *distances, idx_t *labels) const override
Definition: IndexBinaryIVF.cpp:128
faiss::Index
Definition: Index.h:62
faiss::IVFSearchParameters::max_codes
size_t max_codes
max nb of codes to visit to do a query
Definition: IndexIVF.h:63
faiss::index_factory
Index * index_factory(int d, const char *description_in, MetricType metric)
Definition: AutoTune.cpp:722
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:45