Skip to content

F
faiss
Unverified Commit 76bec0b5 authored by Lucas Hosseini's avatar Lucas Hosseini Committed by GitHub
Browse files
Options

Facebook sync (#573) 

Features:

- automatic tracking of C++ references in Python
- non-intel platforms supported -- some functions optimized for ARM
- override nprobe for concurrent searches
- support for floating-point quantizers in binary indexes
Bug fixes:

- no more segfaults in python (I know it's the same as the first feature but it's important!)
- fix GpuIndexIVFFlat issues for float32 with 64 / 128 dims
- fix sharding of flat indexes on GPU with index_cpu_to_gpu_multiple
parent 19cea3d2
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1994 additions and 657 deletions
AutoTune.cpp
View file @ 76bec0b5
......@@ -518,6 +518,21 @@ void ParameterSpace::set_index_parameter (
return;
}
}
if (name == "efSearch") {
if (DC (IndexHNSW)) {
ix->hnsw.efSearch = int(val);
return;
}
if (DC (IndexIVF)) {
if (IndexHNSW *cq =
dynamic_cast<IndexHNSW *>(ix->quantizer)) {
cq->hnsw.efSearch = int(val);
return;
}
}
}
FAISS_THROW_FMT ("ParameterSpace::set_index_parameter:"
"could not set parameter %s",
name.c_str());
......@@ -682,6 +697,7 @@ struct VTChain {
char get_trains_alone(const Index *coarse_quantizer) {
return
dynamic_cast<const MultiIndexQuantizer*>(coarse_quantizer) ? 1 :
dynamic_cast<const IndexHNSWFlat*>(coarse_quantizer) ? 2 :
0;
}
......@@ -738,6 +754,11 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
} else if (stok == "L2norm") {
vt_1 = new NormalizationTransform (d, 2.0);
// coarse quantizers
} else if (!coarse_quantizer &&
sscanf (tok, "IVF%d_HNSW%d", &ncentroids, &M) == 2) {
FAISS_THROW_IF_NOT (metric == METRIC_L2);
coarse_quantizer_1 = new IndexHNSWFlat (d, M);
} else if (!coarse_quantizer &&
sscanf (tok, "IVF%d", &ncentroids) == 1) {
......@@ -935,4 +956,5 @@ IndexBinary *index_binary_factory(int d, const char *description)
}
} // namespace faiss
AutoTune.h
View file @ 76bec0b5
......@@ -205,8 +205,6 @@ Index *index_factory (int d, const char *description,
IndexBinary *index_binary_factory (int d, const char *description);
} // namespace faiss
......
AuxIndexStructures.h
View file @ 76bec0b5
......@@ -198,7 +198,7 @@ struct IOWriter {
struct VectorIOReader:IOReader {
const std::vector<uint8_t> data;
std::vector<uint8_t> data;
size_t rp = 0;
size_t operator()(void *ptr, size_t size, size_t nitems) override;
};
......
IVFlib.cpp 0 → 100644
View file @ 76bec0b5
This diff is collapsed.
IVFlib.h 0 → 100644
View file @ 76bec0b5
/**
* 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.
*/
// -*- c++ -*-
#ifndef FAISS_IVFLIB_H
#define FAISS_IVFLIB_H
/** Since IVF (inverted file) indexes are of so much use for
* large-scale use cases, we group a few functions related to them in
* this small library. Most functions work both on IndexIVFs and
* IndexIVFs embedded within an IndexPreTransform.
*/
#include <vector>
#include "IndexIVF.h"
namespace faiss { namespace ivflib {
/** check if two indexes have the same parameters and are trained in
* the same way, otherwise throw. */
void check_compatible_for_merge (const Index * index1,
const Index * index2);
/** get an IndexIVF from an index. The index may be an IndexIVF or
* some wrapper class that encloses an IndexIVF
*
* throws an exception if this is not the case.
*/
const IndexIVF * extract_index_ivf (const Index * index);
IndexIVF * extract_index_ivf (Index * index);
/** Merge index1 into index0. Works on IndexIVF's and IndexIVF's
* embedded in a IndexPreTransform. On output, the index1 is empty.
*
* @param shift_ids: translate the ids from index1 to index0->prev_ntotal
*/
void merge_into(Index *index0, Index *index1, bool shift_ids);
typedef Index::idx_t idx_t;
/* 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 centroid_ids
* cluster id each object belongs to, size num_objects
*/
void search_centroid(Index *index,
const float* x, int n,
idx_t* 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_return_centroids(Index *index,
size_t n,
const float* xin,
long k,
float *distances,
idx_t* labels,
idx_t* query_centroid_ids,
idx_t* result_centroid_ids);
/** A set of IndexIVFs concatenated together in a FIFO fashion.
* at each "step", the oldest index slice is removed and a new index is added.
*/
struct SlidingIndexWindow {
/// common index that contains the sliding window
Index * index;
/// InvertedLists of index
ArrayInvertedLists *ils;
/// number of slices currently in index
int n_slice;
/// same as index->nlist
size_t nlist;
/// cumulative list sizes at each slice
std::vector<std::vector<size_t> > sizes;
/// index should be initially empty and trained
SlidingIndexWindow (Index *index);
/** Add one index to the current index and remove the oldest one.
*
* @param sub_index slice to swap in (can be NULL)
* @param remove_oldest if true, remove the oldest slices */
void step(const Index *sub_index, bool remove_oldest);
};
/// Get a subset of inverted lists [i0, i1)
ArrayInvertedLists * get_invlist_range (const Index *index,
long i0, long i1);
/// Set a subset of inverted lists
void set_invlist_range (Index *index, long i0, long i1,
ArrayInvertedLists * src);
// search an IndexIVF, possibly embedded in an IndexPreTransform
// with given parameters
void search_with_parameters (const Index *index,
idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
IVFSearchParameters *params);
} } // namespace faiss::ivflib
#endif
IndexBinaryFromFloat.cpp 0 → 100644
View file @ 76bec0b5
/**
* 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.
*/
// -*- c++ -*-
#include "IndexBinaryFromFloat.h"
#include <memory>
#include "utils.h"
namespace faiss {
IndexBinaryFromFloat::IndexBinaryFromFloat(Index *index)
: IndexBinary(index->d),
index(index),
own_fields(false) {
is_trained = index->is_trained;
ntotal = index->ntotal;
}
IndexBinaryFromFloat::~IndexBinaryFromFloat() {
if (own_fields) {
delete index;
}
}
void IndexBinaryFromFloat::add(idx_t n, const uint8_t *x) {
constexpr idx_t bs = 32768;
std::unique_ptr<float[]> xf(new float[bs * d]);
for (idx_t b = 0; b < n; b += bs) {
idx_t bn = std::min(bs, n - b);
binary_to_real(bn * d, x + b * code_size, xf.get());
index->add(bn, xf.get());
}
ntotal = index->ntotal;
}
void IndexBinaryFromFloat::reset() {
index->reset();
ntotal = index->ntotal;
}
void IndexBinaryFromFloat::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels) const {
constexpr idx_t bs = 32768;
std::unique_ptr<float[]> xf(new float[bs * d]);
std::unique_ptr<float[]> df(new float[bs * k]);
for (idx_t b = 0; b < n; b += bs) {
idx_t bn = std::min(bs, n - b);
binary_to_real(bn * d, x + b * code_size, xf.get());
index->search(bn, xf.get(), k, df.get(), labels + b * k);
for (int i = 0; i < bn * k; ++i) {
distances[b * k + i] = int32_t(std::round(df[i] / 4.0));
}
}
}
void IndexBinaryFromFloat::train(idx_t n, const uint8_t *x) {
std::unique_ptr<float[]> xf(new float[n * d]);
binary_to_real(n * d, x, xf.get());
index->train(n, xf.get());
is_trained = true;
ntotal = index->ntotal;
}
} // namespace faiss
IndexBinaryFromFloat.h 0 → 100644
View file @ 76bec0b5
/**
* 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.
*/
// -*- c++ -*-
#ifndef FAISS_INDEX_BINARY_FROM_FLOAT_H
#define FAISS_INDEX_BINARY_FROM_FLOAT_H
#include "IndexBinary.h"
namespace faiss {
class Index;
/** IndexBinary backed by a float Index.
*
* Supports adding vertices and searching them.
*
* All queries are symmetric because there is no distinction between codes and
* vectors.
*/
struct IndexBinaryFromFloat : IndexBinary {
Index *index;
bool own_fields; ///< Whether object owns the index pointer.
explicit IndexBinaryFromFloat(Index *index);
~IndexBinaryFromFloat();
void add(idx_t n, const uint8_t *x) override;
void reset() override;
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels) const override;
void train(idx_t n, const uint8_t *x) override;
};
} // namespace faiss
#endif // FAISS_INDEX_BINARY_FROM_FLOAT_H
IndexBinaryIVF.cpp
View file @ 76bec0b5
......@@ -24,7 +24,6 @@
namespace faiss {
IndexBinaryIVF::IndexBinaryIVF(IndexBinary *quantizer, size_t d, size_t nlist)
: IndexBinary(d),
invlists(new ArrayInvertedLists(nlist, code_size)),
......@@ -291,26 +290,7 @@ void IndexBinaryIVF::merge_from(IndexBinaryIVF &other, idx_t add_id) {
FAISS_THROW_IF_NOT_MSG(typeid (*this) == typeid (other),
"can only merge indexes of the same type");
InvertedLists *oivf = other.invlists;
#pragma omp parallel for
for (long i = 0; i < nlist; i++) {
size_t list_size = oivf->list_size(i);
const idx_t * ids = oivf->get_ids(i);
if (add_id == 0) {
invlists->add_entries(i, list_size, ids,
oivf->get_codes(i));
} else {
std::vector <idx_t> new_ids(list_size);
for (size_t j = 0; j < list_size; j++) {
new_ids [j] = ids[j] + add_id;
}
invlists->add_entries(i, list_size, new_ids.data(),
oivf->get_codes(i));
}
oivf->resize(i, 0);
}
invlists->merge_from (other.invlists, add_id);
ntotal += other.ntotal;
other.ntotal = 0;
......@@ -327,33 +307,6 @@ void IndexBinaryIVF::replace_invlists(InvertedLists *il, bool own) {
}
namespace {
void binary_to_real(int d, const uint8_t *x_in, float *x_out) {
for (int j = 0; j < d; ++j) {
if ((x_in[j / 8] & (1 << (j % 8))) == 0) {
x_out[j] = -1.0;
} else {
x_out[j] = 1.0;
}
}
}
void real_to_binary(int d, const float *x_in, uint8_t *x_out) {
for (int j = 0; j < d; ++j) {
if (x_in[j] > 0) {
x_out[j / 8] |= (1 << (j % 8));
} else {
x_out[j / 8] &= ~(1 << (j % 8));
}
}
}
} // namespace
void IndexBinaryIVF::train_q1(size_t n, const uint8_t *x, bool verbose) {
if (quantizer->is_trained && (quantizer->ntotal == nlist)) {
if (verbose)
......@@ -366,22 +319,15 @@ void IndexBinaryIVF::train_q1(size_t n, const uint8_t *x, bool verbose) {
quantizer->reset();
std::unique_ptr<float[]> x_f(new float[n * d]);
for (int i = 0; i < n; ++i) {
binary_to_real(d,
x + i * code_size,
x_f.get() + i * d);
}
binary_to_real(n * d, x, x_f.get());
IndexFlatL2 index_tmp(d);
clus.train(n, x_f.get(), index_tmp);
std::unique_ptr<uint8_t[]> x_b(new uint8_t[clus.k * code_size]);
for (int i = 0; i < clus.k; ++i) {
real_to_binary(d,
clus.centroids.data() + i * d,
x_b.get() + i * code_size);
}
real_to_binary(d * clus.k, clus.centroids.data(), x_b.get());
quantizer->add(clus.k, x_b.get());
quantizer->is_trained = true;
}
......@@ -426,7 +372,9 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
nlistv++;
size_t list_size = ivf.invlists->list_size(key);
const uint8_t *list_vecs = (const uint8_t*)ivf.invlists->get_codes(key);
InvertedLists::ScopedCodes scodes (ivf.invlists, key);
const uint8_t *list_vecs = scodes.get();
const Index::idx_t *ids = store_pairs
? nullptr
: ivf.invlists->get_ids(key);
......@@ -443,6 +391,9 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
maxheap_push(k, disi, idxi, disij, id);
}
}
if (ids)
ivf.invlists->release_ids (ids);
nscan += list_size;
if (max_codes && nscan >= max_codes)
break;
......@@ -504,7 +455,8 @@ void search_knn_hamming_count(const IndexBinaryIVF& ivf,
nlistv++;
size_t list_size = ivf.invlists->list_size(key);
const uint8_t *list_vecs = (const uint8_t*)ivf.invlists->get_codes(key);
InvertedLists::ScopedCodes scodes (ivf.invlists, key);
const uint8_t *list_vecs = scodes.get();
const Index::idx_t *ids = store_pairs
? nullptr
: ivf.invlists->get_ids(key);
......@@ -515,6 +467,9 @@ void search_knn_hamming_count(const IndexBinaryIVF& ivf,
long id = store_pairs ? (key << 32 | j) : ids[j];
csi.update_counter(yj, id);
}
if (ids)
ivf.invlists->release_ids (ids);
nscan += list_size;
if (max_codes && nscan >= max_codes)
break;
......
IndexHNSW.cpp
View file @ 76bec0b5
......@@ -26,7 +26,9 @@
#include <unistd.h>
#include <stdint.h>
#ifdef __SSE__
#include <immintrin.h>
#endif
#include "utils.h"
#include "Heap.h"
......@@ -1869,6 +1871,7 @@ struct DistanceXPQ4: Distance2Level {
float operator () (storage_idx_t i) override
{
#ifdef __SSE__
const uint8_t *code = storage.codes.data() + i * storage.code_size;
long key = 0;
memcpy (&key, code, storage.code_size_1);
......@@ -1892,6 +1895,9 @@ struct DistanceXPQ4: Distance2Level {
accu = _mm_hadd_ps (accu, accu);
accu = _mm_hadd_ps (accu, accu);
return _mm_cvtss_f32 (accu);
#else
FAISS_THROW_MSG("not implemented for non-x64 platforms");
#endif
}
};
......@@ -1920,6 +1926,7 @@ struct Distance2xXPQ4: Distance2Level {
long key01 = 0;
memcpy (&key01, code, storage.code_size_1);
code += storage.code_size_1;
#ifdef __SSE__
// walking pointers
const float *qa = q;
......@@ -1945,6 +1952,9 @@ struct Distance2xXPQ4: Distance2Level {
accu = _mm_hadd_ps (accu, accu);
accu = _mm_hadd_ps (accu, accu);
return _mm_cvtss_f32 (accu);
#else
FAISS_THROW_MSG("not implemented for non-x64 platforms");
#endif
}
};
......@@ -1957,6 +1967,7 @@ HNSW::DistanceComputer * IndexHNSW2Level::get_distance_computer () const
dynamic_cast<Index2Layer*>(storage);
if (storage2l) {
#ifdef __SSE__
const MultiIndexQuantizer *mi =
dynamic_cast<MultiIndexQuantizer*> (storage2l->q1.quantizer);
......@@ -1971,6 +1982,7 @@ HNSW::DistanceComputer * IndexHNSW2Level::get_distance_computer () const
if (fl && storage2l->pq.dsub == 4) {
return new DistanceXPQ4(*storage2l);
}
#endif
}
// IVFPQ and cases not handled above
......
IndexIVF.cpp
View file @ 76bec0b5
......@@ -21,6 +21,9 @@
namespace faiss {
using ScopedIds = InvertedLists::ScopedIds;
using ScopedCodes = InvertedLists::ScopedCodes;
/*****************************************
* Level1Quantizer implementation
******************************************/
......@@ -98,120 +101,6 @@ void Level1Quantizer::train_q1 (size_t n, const float *x, bool verbose, MetricTy
}
}
/*****************************************
* InvertedLists implementation
******************************************/
InvertedLists::InvertedLists (size_t nlist, size_t code_size):
nlist (nlist), code_size (code_size)
{
}
InvertedLists::~InvertedLists ()
{}
InvertedLists::idx_t InvertedLists::get_single_id (
size_t list_no, size_t offset) const
{
assert (offset < list_size (list_no));
return get_ids(list_no)[offset];
}
void InvertedLists::prefetch_lists (const long *, int) const
{}
const uint8_t * InvertedLists::get_single_code (
size_t list_no, size_t offset) const
{
assert (offset < list_size (list_no));
return get_codes(list_no) + offset * code_size;
}
size_t InvertedLists::add_entry (size_t list_no, idx_t theid,
const uint8_t *code)
{
return add_entries (list_no, 1, &theid, code);
}
void InvertedLists::update_entry (size_t list_no, size_t offset,
idx_t id, const uint8_t *code)
{
update_entries (list_no, offset, 1, &id, code);
}
void InvertedLists::reset () {
for (size_t i = 0; i < nlist; i++) {
resize (i, 0);
}
}
/*****************************************
* ArrayInvertedLists implementation
******************************************/
ArrayInvertedLists::ArrayInvertedLists (size_t nlist, size_t code_size):
InvertedLists (nlist, code_size)
{
ids.resize (nlist);
codes.resize (nlist);
}
size_t ArrayInvertedLists::add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids_in, const uint8_t *code)
{
if (n_entry == 0) return 0;
assert (list_no < nlist);
size_t o = ids [list_no].size();
ids [list_no].resize (o + n_entry);
memcpy (&ids[list_no][o], ids_in, sizeof (ids_in[0]) * n_entry);
codes [list_no].resize ((o + n_entry) * code_size);
memcpy (&codes[list_no][o * code_size], code, code_size * n_entry);
return o;
}
size_t ArrayInvertedLists::list_size(size_t list_no) const
{
assert (list_no < nlist);
return ids[list_no].size();
}
const uint8_t * ArrayInvertedLists::get_codes (size_t list_no) const
{
assert (list_no < nlist);
return codes[list_no].data();
}
const InvertedLists::idx_t * ArrayInvertedLists::get_ids (size_t list_no) const
{
assert (list_no < nlist);
return ids[list_no].data();
}
void ArrayInvertedLists::resize (size_t list_no, size_t new_size)
{
ids[list_no].resize (new_size);
codes[list_no].resize (new_size * code_size);
}
void ArrayInvertedLists::update_entries (
size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids_in, const uint8_t *codes_in)
{
assert (list_no < nlist);
assert (n_entry + offset <= ids[list_no].size());
memcpy (&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry);
memcpy (&codes[list_no][offset * code_size], codes_in, code_size * n_entry);
}
ArrayInvertedLists::~ArrayInvertedLists ()
{}
/*****************************************
......@@ -262,7 +151,7 @@ void IndexIVF::make_direct_map (bool new_maintain_direct_map)
direct_map.resize (ntotal, -1);
for (size_t key = 0; key < nlist; key++) {
size_t list_size = invlists->list_size (key);
const idx_t *idlist = invlists->get_ids (key);
ScopedIds idlist (invlists, key);
for (long ofs = 0; ofs < list_size; ofs++) {
FAISS_THROW_IF_NOT_MSG (
......@@ -312,12 +201,12 @@ void IndexIVF::reconstruct_n (idx_t i0, idx_t ni, float* recons) const
for (long list_no = 0; list_no < nlist; list_no++) {
size_t list_size = invlists->list_size (list_no);
const Index::idx_t * idlist = invlists->get_ids (list_no);
ScopedIds idlist (invlists, list_no);
for (long offset = 0; offset < list_size; offset++) {
long id = idlist[offset];
if (!(id >= i0 && id < i0 + ni)) {
continue;
continue;
}
float* reconstructed = recons + (id - i0) * d;
......@@ -390,14 +279,14 @@ long IndexIVF::remove_ids (const IDSelector & sel)
#pragma omp parallel for
for (long i = 0; i < nlist; i++) {
long l0 = invlists->list_size (i), l = l0, j = 0;
const idx_t *idsi = invlists->get_ids (i);
ScopedIds idsi (invlists, i);
while (j < l) {
if (sel.is_member (idsi[j])) {
l--;
invlists->update_entry (
i, j,
invlists->get_single_id (i, l),
invlists->get_single_code (i, l));
ScopedCodes (invlists, i, l).get());
} else {
j++;
}
......@@ -472,38 +361,26 @@ void IndexIVF::print_stats () const
}
void IndexIVF::merge_from (IndexIVF &other, idx_t add_id)
void IndexIVF::check_compatible_for_merge (const IndexIVF &other) const
{
// minimal sanity checks
FAISS_THROW_IF_NOT (other.d == d);
FAISS_THROW_IF_NOT (other.nlist == nlist);
FAISS_THROW_IF_NOT (other.code_size == code_size);
FAISS_THROW_IF_NOT_MSG ((!maintain_direct_map &&
!other.maintain_direct_map),
"direct map copy not implemented");
FAISS_THROW_IF_NOT_MSG (typeid (*this) == typeid (other),
"can only merge indexes of the same type");
}
InvertedLists *oivf = other.invlists;
#pragma omp parallel for
for (long i = 0; i < nlist; i++) {
size_t list_size = oivf->list_size (i);
const idx_t * ids = oivf->get_ids (i);
if (add_id == 0) {
invlists->add_entries (i, list_size, ids,
oivf->get_codes (i));
} else {
std::vector <idx_t> new_ids (list_size);
for (size_t j = 0; j < list_size; j++) {
new_ids [j] = ids[j] + add_id;
}
void IndexIVF::merge_from (IndexIVF &other, idx_t add_id)
{
check_compatible_for_merge (other);
FAISS_THROW_IF_NOT_MSG ((!maintain_direct_map &&
!other.maintain_direct_map),
"direct map copy not implemented");
invlists->add_entries (i, list_size, new_ids.data(),
oivf->get_codes (i));
}
oivf->resize (i, 0);
}
invlists->merge_from (other.invlists, add_id);
ntotal += other.ntotal;
other.ntotal = 0;
......@@ -542,7 +419,7 @@ void IndexIVF::copy_subset_to (IndexIVF & other, int subset_type,
for (long list_no = 0; list_no < nlist; list_no++) {
size_t n = invlists->list_size (list_no);
const idx_t *ids_in = invlists->get_ids (list_no);
ScopedIds ids_in (invlists, list_no);
if (subset_type == 0) {
for (long i = 0; i < n; i++) {
......@@ -550,7 +427,7 @@ void IndexIVF::copy_subset_to (IndexIVF & other, int subset_type,
if (a1 <= id && id < a2) {
oivf->add_entry (list_no,
invlists->get_single_id (list_no, i),
invlists->get_single_code (list_no, i));
ScopedCodes (invlists, list_no, i).get());
other.ntotal++;
}
}
......@@ -560,7 +437,7 @@ void IndexIVF::copy_subset_to (IndexIVF & other, int subset_type,
if (id % a1 == a2) {
oivf->add_entry (list_no,
invlists->get_single_id (list_no, i),
invlists->get_single_code (list_no, i));
ScopedCodes (invlists, list_no, i).get());
other.ntotal++;
}
}
......@@ -575,7 +452,7 @@ void IndexIVF::copy_subset_to (IndexIVF & other, int subset_type,
for (long i = i1; i < i2; i++) {
oivf->add_entry (list_no,
invlists->get_single_id (list_no, i),
invlists->get_single_code (list_no, i));
ScopedCodes (invlists, list_no, i).get());
}
other.ntotal += i2 - i1;
......
IndexIVF.h
View file @ 76bec0b5
......@@ -16,6 +16,7 @@
#include "Index.h"
#include "InvertedLists.h"
#include "Clustering.h"
#include "Heap.h"
......@@ -56,91 +57,6 @@ struct Level1Quantizer {
};
/** Table of inverted lists
* multithreading rules:
* - concurrent read accesses are allowed
* - concurrent update accesses are allowed
* - for resize and add_entries, only concurrent access to different lists
* are allowed
*/
struct InvertedLists {
typedef Index::idx_t idx_t;
size_t nlist; ///< number of possible key values
size_t code_size; ///< code size per vector in bytes
InvertedLists (size_t nlist, size_t code_size);
/*************************
* Read only functions */
/// get the size of a list
virtual size_t list_size(size_t list_no) const = 0;
/// @return codes size list_size * code_size
virtual const uint8_t * get_codes (size_t list_no) const = 0;
/// @return ids size list_size
virtual const idx_t * get_ids (size_t list_no) const = 0;
/// @return a single id in an inverted list
virtual idx_t get_single_id (size_t list_no, size_t offset) const;
/// @return a single code in an inverted list
virtual const uint8_t * get_single_code (
size_t list_no, size_t offset) const;
/// prepare the following lists (default does nothing)
/// a list can be -1 hence the signed long
virtual void prefetch_lists (const long *list_nos, int nlist) const;
/*************************
* writing functions */
/// add one entry to an inverted list
virtual size_t add_entry (size_t list_no, idx_t theid,
const uint8_t *code);
virtual size_t add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids, const uint8_t *code) = 0;
virtual void update_entry (size_t list_no, size_t offset,
idx_t id, const uint8_t *code);
virtual void update_entries (size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids, const uint8_t *code) = 0;
virtual void resize (size_t list_no, size_t new_size) = 0;
virtual void reset ();
virtual ~InvertedLists ();
};
struct ArrayInvertedLists: InvertedLists {
std::vector < std::vector<uint8_t> > codes; // binary codes, size nlist
std::vector < std::vector<idx_t> > ids; ///< Inverted lists for indexes
ArrayInvertedLists (size_t nlist, size_t code_size);
size_t list_size(size_t list_no) const override;
const uint8_t * get_codes (size_t list_no) const override;
const idx_t * get_ids (size_t list_no) const override;
size_t add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids, const uint8_t *code) override;
void update_entries (size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids, const uint8_t *code) override;
void resize (size_t list_no, size_t new_size) override;
virtual ~ArrayInvertedLists ();
};
struct IVFSearchParameters {
size_t nprobe; ///< number of probes at query time
......@@ -273,6 +189,11 @@ struct IndexIVF: Index, Level1Quantizer {
long remove_ids(const IDSelector& sel) override;
/** check that the two indexes are compatible (ie, they are
* trained in the same way and have the same
* parameters). Otherwise throw. */
void check_compatible_for_merge (const IndexIVF &other) const;
/** moves the entries from another dataset to self. On output,
* other is empty. add_id is added to all moved ids (for
* sequential ids, this would be this->ntotal */
......
IndexIVFFlat.cpp
View file @ 76bec0b5
......@@ -18,6 +18,7 @@
#include "IndexFlat.h"
#include "AuxIndexStructures.h"
namespace faiss {
......@@ -33,8 +34,6 @@ IndexIVFFlat::IndexIVFFlat (Index * quantizer,
}
void IndexIVFFlat::add_with_ids (idx_t n, const float * x, const long *xids)
{
add_core (n, x, xids, nullptr);
......@@ -122,8 +121,8 @@ void search_knn_for_ivf (const IndexIVFFlat & ivf,
nlistv++;
size_t list_size = ivf.invlists->list_size(key);
const float * list_vecs =
(const float*)ivf.invlists->get_codes (key);
InvertedLists::ScopedCodes scodes (ivf.invlists, key);
const float * list_vecs = (const float*)scodes.get();
const Index::idx_t * ids = store_pairs ? nullptr :
ivf.invlists->get_ids (key);
......@@ -137,6 +136,10 @@ void search_knn_for_ivf (const IndexIVFFlat & ivf,
heap_push<C> (k, simi, idxi, dis, id);
}
}
if (ids) {
ivf.invlists->release_ids (ids);
}
nscan += list_size;
if (max_codes && nscan >= max_codes)
break;
......@@ -213,9 +216,9 @@ void IndexIVFFlat::range_search (idx_t nx, const float *x, float radius,
}
const size_t list_size = invlists->list_size(key);
const float * list_vecs =
(const float*)invlists->get_codes (key);
const Index::idx_t * ids = invlists->get_ids (key);
InvertedLists::ScopedCodes scodes (invlists, key);
const float * list_vecs = (const float*)scodes.get();
InvertedLists::ScopedIds ids (invlists, key);
for (size_t j = 0; j < list_size; j++) {
const float * yj = list_vecs + d * j;
......@@ -355,11 +358,12 @@ void IndexIVFFlatDedup::add_with_ids(
const float *xi = x + i * d;
// search if there is already an entry with that id
const uint8_t * codes = invlists->get_codes (list_no);
InvertedLists::ScopedCodes codes (invlists, list_no);
long n = invlists->list_size (list_no);
long offset = -1;
for (long o = 0; o < n; o++) {
if (!memcmp (codes + o * code_size,
if (!memcmp (codes.get() + o * code_size,
xi, code_size)) {
offset = o;
break;
......@@ -479,7 +483,7 @@ long IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
#pragma omp parallel for
for (long i = 0; i < nlist; i++) {
long l0 = invlists->list_size (i), l = l0, j = 0;
const idx_t *idsi = invlists->get_ids (i);
InvertedLists::ScopedIds idsi (invlists, i);
while (j < l) {
if (sel.is_member (idsi[j])) {
if (replace.count(idsi[j]) == 0) {
......@@ -487,12 +491,12 @@ long IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
invlists->update_entry (
i, j,
invlists->get_single_id (i, l),
invlists->get_single_code (i, l));
InvertedLists::ScopedCodes (invlists, i, l).get());
} else {
invlists->update_entry (
i, j,
replace[idsi[j]],
invlists->get_single_code (i, j));
InvertedLists::ScopedCodes (invlists, i, j).get());
j++;
}
} else {
......
IndexIVFPQ.cpp
View file @ 76bec0b5
......@@ -421,11 +421,15 @@ void IndexIVFPQ::precompute_table ()
namespace {
static uint64_t get_cycles () {
#ifdef __x86_64__
uint32_t high, low;
asm volatile("rdtsc \n\t"
: "=a" (low),
"=d" (high));
return ((uint64_t)high << 32) | (low);
#else
return 0;
#endif
}
#define TIC t0 = get_cycles()
......@@ -987,8 +991,9 @@ void IndexIVFPQ::search_preassigned (idx_t nx, const float *qx, idx_t k,
if (list_size == 0) continue;
qt.init_list (key, coarse_dis_i[ik],
list_size, invlists->get_ids (key),
invlists->get_codes (key));
list_size,
InvertedLists::ScopedIds (invlists, key).get(),
InvertedLists::ScopedCodes (invlists, key).get());
TIC;
if (polysemous_ht > 0) {
......@@ -1063,10 +1068,11 @@ size_t IndexIVFPQ::find_duplicates (idx_t *dup_ids, size_t *lims) const
size_t n = invlists->list_size (list_no);
std::vector<int> ord (n);
for (int i = 0; i < n; i++) ord[i] = i;
CodeCmp cs = { invlists->get_codes (list_no), code_size };
InvertedLists::ScopedCodes codes (invlists, list_no);
CodeCmp cs = { codes.get(), code_size };
std::sort (ord.begin(), ord.end(), cs);
const idx_t *list_ids = invlists->get_ids (list_no);
InvertedLists::ScopedIds list_ids (invlists, list_no);
int prev = -1; // all elements from prev to i-1 are equal
for (int i = 0; i < n; i++) {
if (prev >= 0 && cs.cmp (ord [prev], ord [i]) == 0) {
......
IndexScalarQuantizer.cpp
View file @ 76bec0b5
......@@ -15,7 +15,9 @@
#include <omp.h>
#ifdef __SSE__
#include <immintrin.h>
#endif
#include "utils.h"
......@@ -143,11 +145,13 @@ float decode_fp16 (uint16_t x) {
// https://github.com/ispc/ispc/blob/master/stdlib.ispc
float floatbits (uint32_t x) {
return *(float*)&x;
void *xptr = &x;
return *(float*)xptr;
}
uint32_t intbits (float f) {
return *(uint32_t*)&f;
void *fptr = &f;
return *(uint32_t*)fptr;
}
......@@ -1179,7 +1183,8 @@ void search_with_probes_ip (const IndexIVFScalarQuantizer & index,
float accu0 = cent_dis[i];
const size_t list_size = index.invlists->list_size (list_no);
const uint8_t * codes = index.invlists->get_codes (list_no);
InvertedLists::ScopedCodes scodes (index.invlists, list_no);
const uint8_t *codes = scodes.get();
const idx_t * ids =
store_pairs ? nullptr : index.invlists->get_ids (list_no);
......@@ -1196,6 +1201,9 @@ void search_with_probes_ip (const IndexIVFScalarQuantizer & index,
}
codes += code_size;
}
if (ids) {
index.invlists->release_ids (ids);
}
nscan += list_size;
if (max_codes && nscan > max_codes)
break;
......@@ -1225,7 +1233,8 @@ void search_with_probes_L2 (const IndexIVFScalarQuantizer & index,
if (list_no < 0) break;
const size_t list_size = index.invlists->list_size (list_no);
const uint8_t * codes = index.invlists->get_codes (list_no);
InvertedLists::ScopedCodes scodes (index.invlists, list_no);
const uint8_t *codes = scodes.get();
const idx_t * ids =
store_pairs ? nullptr : index.invlists->get_ids (list_no);
......@@ -1243,6 +1252,9 @@ void search_with_probes_L2 (const IndexIVFScalarQuantizer & index,
}
codes += code_size;
}
if (ids) {
index.invlists->release_ids (ids);
}
nscan += list_size;
if (max_codes && nscan > max_codes)
break;
......
InvertedLists.cpp 0 → 100644
View file @ 76bec0b5
/**
* 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.
*/
// -*- c++ -*-
#include "InvertedLists.h"
#include <cstdio>
#include "utils.h"
#include "FaissAssert.h"
namespace faiss {
using ScopedIds = InvertedLists::ScopedIds;
using ScopedCodes = InvertedLists::ScopedCodes;
/*****************************************
* InvertedLists implementation
******************************************/
InvertedLists::InvertedLists (size_t nlist, size_t code_size):
nlist (nlist), code_size (code_size)
{
}
InvertedLists::~InvertedLists ()
{}
InvertedLists::idx_t InvertedLists::get_single_id (
size_t list_no, size_t offset) const
{
assert (offset < list_size (list_no));
return get_ids(list_no)[offset];
}
void InvertedLists::release_codes (const uint8_t *) const
{}
void InvertedLists::release_ids (const idx_t *) const
{}
void InvertedLists::prefetch_lists (const long *, int) const
{}
const uint8_t * InvertedLists::get_single_code (
size_t list_no, size_t offset) const
{
assert (offset < list_size (list_no));
return get_codes(list_no) + offset * code_size;
}
size_t InvertedLists::add_entry (size_t list_no, idx_t theid,
const uint8_t *code)
{
return add_entries (list_no, 1, &theid, code);
}
void InvertedLists::update_entry (size_t list_no, size_t offset,
idx_t id, const uint8_t *code)
{
update_entries (list_no, offset, 1, &id, code);
}
void InvertedLists::reset () {
for (size_t i = 0; i < nlist; i++) {
resize (i, 0);
}
}
void InvertedLists::merge_from (InvertedLists *oivf, size_t add_id) {
#pragma omp parallel for
for (long i = 0; i < nlist; i++) {
size_t list_size = oivf->list_size (i);
ScopedIds ids (oivf, i);
if (add_id == 0) {
add_entries (i, list_size, ids.get (),
ScopedCodes (oivf, i).get());
} else {
std::vector <idx_t> new_ids (list_size);
for (size_t j = 0; j < list_size; j++) {
new_ids [j] = ids[j] + add_id;
}
add_entries (i, list_size, new_ids.data(),
ScopedCodes (oivf, i).get());
}
oivf->resize (i, 0);
}
}
/*****************************************
* ArrayInvertedLists implementation
******************************************/
ArrayInvertedLists::ArrayInvertedLists (size_t nlist, size_t code_size):
InvertedLists (nlist, code_size)
{
ids.resize (nlist);
codes.resize (nlist);
}
size_t ArrayInvertedLists::add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids_in, const uint8_t *code)
{
if (n_entry == 0) return 0;
assert (list_no < nlist);
size_t o = ids [list_no].size();
ids [list_no].resize (o + n_entry);
memcpy (&ids[list_no][o], ids_in, sizeof (ids_in[0]) * n_entry);
codes [list_no].resize ((o + n_entry) * code_size);
memcpy (&codes[list_no][o * code_size], code, code_size * n_entry);
return o;
}
size_t ArrayInvertedLists::list_size(size_t list_no) const
{
assert (list_no < nlist);
return ids[list_no].size();
}
const uint8_t * ArrayInvertedLists::get_codes (size_t list_no) const
{
assert (list_no < nlist);
return codes[list_no].data();
}
const InvertedLists::idx_t * ArrayInvertedLists::get_ids (size_t list_no) const
{
assert (list_no < nlist);
return ids[list_no].data();
}
void ArrayInvertedLists::resize (size_t list_no, size_t new_size)
{
ids[list_no].resize (new_size);
codes[list_no].resize (new_size * code_size);
}
void ArrayInvertedLists::update_entries (
size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids_in, const uint8_t *codes_in)
{
assert (list_no < nlist);
assert (n_entry + offset <= ids[list_no].size());
memcpy (&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry);
memcpy (&codes[list_no][offset * code_size], codes_in, code_size * n_entry);
}
ArrayInvertedLists::~ArrayInvertedLists ()
{}
/*****************************************
* ConcatenatedInvertedLists implementation
******************************************/
ConcatenatedInvertedLists::ConcatenatedInvertedLists (
int nil, const InvertedLists **ils_in):
InvertedLists (nil > 0 ? ils_in[0]->nlist : 0,
nil > 0 ? ils_in[0]->code_size : 0)
{
FAISS_THROW_IF_NOT (nil > 0);
for (int i = 0; i < nil; i++) {
ils.push_back (ils_in[i]);
FAISS_THROW_IF_NOT (ils_in[i]->code_size == code_size &&
ils_in[i]->nlist == nlist);
}
}
size_t ConcatenatedInvertedLists::list_size(size_t list_no) const
{
size_t sz = 0;
for (int i = 0; i < ils.size(); i++) {
const InvertedLists *il = ils[i];
sz += il->list_size (list_no);
}
return sz;
}
const uint8_t * ConcatenatedInvertedLists::get_codes (size_t list_no) const
{
uint8_t *codes = new uint8_t [code_size * list_size(list_no)], *c = codes;
for (int i = 0; i < ils.size(); i++) {
const InvertedLists *il = ils[i];
size_t sz = il->list_size(list_no) * code_size;
if (sz > 0) {
memcpy (c, ScopedCodes (il, list_no).get(), sz);
c += sz;
}
}
return codes;
}
const uint8_t * ConcatenatedInvertedLists::get_single_code (
size_t list_no, size_t offset) const
{
for (int i = 0; i < ils.size(); i++) {
const InvertedLists *il = ils[i];
size_t sz = il->list_size (list_no);
if (offset < sz) {
// here we have to copy the code, otherwise it will crash at dealloc
uint8_t * code = new uint8_t [code_size];
memcpy (code, ScopedCodes (il, list_no, offset).get(), code_size);
return code;
}
offset -= sz;
}
FAISS_THROW_FMT ("offset %ld unknown", offset);
}
void ConcatenatedInvertedLists::release_codes (const uint8_t *codes) const {
delete [] codes;
}
const Index::idx_t * ConcatenatedInvertedLists::get_ids (size_t list_no) const
{
idx_t *ids = new idx_t [list_size(list_no)], *c = ids;
for (int i = 0; i < ils.size(); i++) {
const InvertedLists *il = ils[i];
size_t sz = il->list_size(list_no);
if (sz > 0) {
memcpy (c, ScopedIds (il, list_no).get(), sz * sizeof(idx_t));
c += sz;
}
}
return ids;
}
Index::idx_t ConcatenatedInvertedLists::get_single_id (
size_t list_no, size_t offset) const
{
for (int i = 0; i < ils.size(); i++) {
const InvertedLists *il = ils[i];
size_t sz = il->list_size (list_no);
if (offset < sz) {
return il->get_single_id (list_no, offset);
}
offset -= sz;
}
FAISS_THROW_FMT ("offset %ld unknown", offset);
}
void ConcatenatedInvertedLists::release_ids (const idx_t *ids) const {
delete [] ids;
}
size_t ConcatenatedInvertedLists::add_entries (
size_t , size_t ,
const idx_t* , const uint8_t *)
{
FAISS_THROW_MSG ("not implemented");
}
void ConcatenatedInvertedLists::update_entries (size_t, size_t , size_t ,
const idx_t *, const uint8_t *)
{
FAISS_THROW_MSG ("not implemented");
}
void ConcatenatedInvertedLists::resize (size_t , size_t )
{
FAISS_THROW_MSG ("not implemented");
}
} // namespace faiss
InvertedLists.h 0 → 100644
View file @ 76bec0b5
/**
* 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.
*/
// -*- c++ -*-
#ifndef FAISS_INVERTEDLISTS_IVF_H
#define FAISS_INVERTEDLISTS_IVF_H
/**
* Definition of inverted lists + a few common classes that implement
* the interface.
*/
#include <vector>
#include "Index.h"
namespace faiss {
/** Table of inverted lists
* multithreading rules:
* - concurrent read accesses are allowed
* - concurrent update accesses are allowed
* - for resize and add_entries, only concurrent access to different lists
* are allowed
*/
struct InvertedLists {
typedef Index::idx_t idx_t;
size_t nlist; ///< number of possible key values
size_t code_size; ///< code size per vector in bytes
InvertedLists (size_t nlist, size_t code_size);
/*************************
* Read only functions */
/// get the size of a list
virtual size_t list_size(size_t list_no) const = 0;
/** get the codes for an inverted list
* must be released by release_codes
*
* @return codes size list_size * code_size
*/
virtual const uint8_t * get_codes (size_t list_no) const = 0;
/** get the ids for an inverted list
* must be released by release_ids
*
* @return ids size list_size
*/
virtual const idx_t * get_ids (size_t list_no) const = 0;
/// release codes returned by get_codes (default implementation is nop
virtual void release_codes (const uint8_t *codes) const;
/// release ids returned by get_ids
virtual void release_ids (const idx_t *ids) const;
/// @return a single id in an inverted list
virtual idx_t get_single_id (size_t list_no, size_t offset) const;
/// @return a single code in an inverted list
/// (should be deallocated with release_codes)
virtual const uint8_t * get_single_code (
size_t list_no, size_t offset) const;
/// prepare the following lists (default does nothing)
/// a list can be -1 hence the signed long
virtual void prefetch_lists (const long *list_nos, int nlist) const;
/*************************
* writing functions */
/// add one entry to an inverted list
virtual size_t add_entry (size_t list_no, idx_t theid,
const uint8_t *code);
virtual size_t add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids, const uint8_t *code) = 0;
virtual void update_entry (size_t list_no, size_t offset,
idx_t id, const uint8_t *code);
virtual void update_entries (size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids, const uint8_t *code) = 0;
virtual void resize (size_t list_no, size_t new_size) = 0;
virtual void reset ();
/// move all entries from oivf (empty on output)
void merge_from (InvertedLists *oivf, size_t add_id);
virtual ~InvertedLists ();
/**************************************
* Scoped inverted lists (for automatic deallocation)
*
* instead of writing:
*
* uint8_t * codes = invlists->get_codes (10);
* ... use codes
* invlists->release_codes(codes)
*
* write:
*
* ScopedCodes codes (invlists, 10);
* ... use codes.get()
* // release called automatically when codes goes out of scope
*
* the following function call also works:
*
* foo (123, ScopedCodes (invlists, 10).get(), 456);
*
*/
struct ScopedIds {
const InvertedLists *il;
const idx_t *ids;
ScopedIds (const InvertedLists *il, size_t list_no):
il (il), ids (il->get_ids (list_no))
{}
const idx_t *get() {return ids; }
idx_t operator [] (size_t i) const {
return ids[i];
}
~ScopedIds () {
il->release_ids (ids);
}
};
struct ScopedCodes {
const InvertedLists *il;
const uint8_t *codes;
ScopedCodes (const InvertedLists *il, size_t list_no):
il (il), codes (il->get_codes (list_no))
{}
ScopedCodes (const InvertedLists *il, size_t list_no, size_t offset):
il (il), codes (il->get_single_code (list_no, offset))
{}
const uint8_t *get() {return codes; }
~ScopedCodes () {
il->release_codes (codes);
}
};
};
/// simple (default) implementation as an array of inverted lists
struct ArrayInvertedLists: InvertedLists {
std::vector < std::vector<uint8_t> > codes; // binary codes, size nlist
std::vector < std::vector<idx_t> > ids; ///< Inverted lists for indexes
ArrayInvertedLists (size_t nlist, size_t code_size);
size_t list_size(size_t list_no) const override;
const uint8_t * get_codes (size_t list_no) const override;
const idx_t * get_ids (size_t list_no) const override;
size_t add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids, const uint8_t *code) override;
void update_entries (size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids, const uint8_t *code) override;
void resize (size_t list_no, size_t new_size) override;
virtual ~ArrayInvertedLists ();
};
/// inverted lists built as the concatenation of a set of invlists
/// (read-only)
struct ConcatenatedInvertedLists: InvertedLists {
std::vector<const InvertedLists *>ils;
/// build InvertedLists by concatenating nil of them
ConcatenatedInvertedLists (int nil, const InvertedLists **ils);
size_t list_size(size_t list_no) const override;
const uint8_t * get_codes (size_t list_no) const override;
const idx_t * get_ids (size_t list_no) const override;
void release_codes (const uint8_t *codes) const override;
void release_ids (const idx_t *ids) const override;
idx_t get_single_id (size_t list_no, size_t offset) const override;
const uint8_t * get_single_code (
size_t list_no, size_t offset) const override;
size_t add_entries (
size_t list_no, size_t n_entry,
const idx_t* ids, const uint8_t *code) override;
void update_entries (size_t list_no, size_t offset, size_t n_entry,
const idx_t *ids, const uint8_t *code) override;
void resize (size_t list_no, size_t new_size) override;
};
} // namespace faiss
#endif
VectorTransform.h
View file @ 76bec0b5
......@@ -134,7 +134,7 @@ struct PCAMatrix: LinearTransform {
* eigenvalues^eigen_power
*
* =0: no whitening
* =-2: full whitening
* =-0.5: full whitening
*/
float eigen_power;
......
benchs/bench_gpu_1bn.py
View file @ 76bec0b5
......@@ -551,9 +551,15 @@ def compute_populated_index(preproc):
print "Aggregate indexes to CPU"
t0 = time.time()
for i in range(ngpu):
index_src = faiss.index_gpu_to_cpu(gpu_index.at(i))
print " index %d size %d" % (i, index_src.ntotal)
if hasattr(gpu_index, 'at'):
# it is a sharded index
for i in range(ngpu):
index_src = faiss.index_gpu_to_cpu(gpu_index.at(i))
print " index %d size %d" % (i, index_src.ntotal)
index_src.copy_subset_to(indexall, 0, 0, nb)
else:
# simple index
index_src = faiss.index_gpu_to_cpu(gpu_index)
index_src.copy_subset_to(indexall, 0, 0, nb)
print " done in %.3f s" % (time.time() - t0)
......
benchs/bench_vector_ops.py
View file @ 76bec0b5
......@@ -10,6 +10,16 @@ import numpy as np
import faiss
import time
swig_ptr = faiss.swig_ptr
if False:
a = np.arange(10, 14).astype('float32')
b = np.arange(20, 24).astype('float32')
faiss.fvec_inner_product (swig_ptr(a), swig_ptr(b), 4)
1/0
xd = 100
yd = 1000000
......@@ -29,9 +39,9 @@ for d in 3, 4, 12, 36, 64:
t0 = time.time()
for i in xrange(xd):
faiss.fvec_inner_products_ny(faiss.swig_ptr(distances[i]),
faiss.swig_ptr(x[i]),
faiss.swig_ptr(y),
faiss.fvec_inner_products_ny(swig_ptr(distances[i]),
swig_ptr(x[i]),
swig_ptr(y),
d, yd)
t1 = time.time()
......@@ -57,9 +67,9 @@ for d in 3, 4, 12, 36, 64:
t0 = time.time()
for i in xrange(xd):
faiss.fvec_L2sqr_ny(faiss.swig_ptr(distances[i]),
faiss.swig_ptr(x[i]),
faiss.swig_ptr(y),
faiss.fvec_L2sqr_ny(swig_ptr(distances[i]),
swig_ptr(x[i]),
swig_ptr(y),
d, yd)
t1 = time.time()
......
gpu/GpuAutoTune.cpp
View file @ 76bec0b5
......@@ -220,6 +220,74 @@ struct ToGpuClonerMultiple: faiss::Cloner, GpuMultipleClonerOptions {
}
Index * clone_Index_to_shards (const Index *index) {
long n = sub_cloners.size();
auto index_ivfpq =
dynamic_cast<const faiss::IndexIVFPQ *>(index);
auto index_ivfflat =
dynamic_cast<const faiss::IndexIVFFlat *>(index);
auto index_flat =
dynamic_cast<const faiss::IndexFlat *>(index);
FAISS_THROW_IF_NOT_MSG (
index_ivfpq || index_ivfflat || index_flat,
"IndexShards implemented only for "
"IndexIVFFlat, IndexFlat and IndexIVFPQ");
std::vector<faiss::Index*> shards(n);
for(long i = 0; i < n; i++) {
// make a shallow copy
if(reserveVecs)
sub_cloners[i].reserveVecs =
(reserveVecs + n - 1) / n;
if (index_ivfpq) {
faiss::IndexIVFPQ idx2(
index_ivfpq->quantizer, index_ivfpq->d,
index_ivfpq->nlist, index_ivfpq->code_size,
index_ivfpq->pq.nbits);
idx2.metric_type = index_ivfpq->metric_type;
idx2.pq = index_ivfpq->pq;
idx2.nprobe = index_ivfpq->nprobe;
idx2.use_precomputed_table = 0;
idx2.is_trained = index->is_trained;
copy_ivf_shard (index_ivfpq, &idx2, n, i);
shards[i] = sub_cloners[i].clone_Index(&idx2);
} else if (index_ivfflat) {
faiss::IndexIVFFlat idx2(
index_ivfflat->quantizer, index->d,
index_ivfflat->nlist, index_ivfflat->metric_type);
idx2.nprobe = index_ivfflat->nprobe;
copy_ivf_shard (index_ivfflat, &idx2, n, i);
shards[i] = sub_cloners[i].clone_Index(&idx2);
} else if (index_flat) {
faiss::IndexFlat idx2 (
index->d, index->metric_type);
shards[i] = sub_cloners[i].clone_Index(&idx2);
if (index->ntotal > 0) {
long i0 = index->ntotal * i / n;
long i1 = index->ntotal * (i + 1) / n;
shards[i]->add (
i1 - i0,
index_flat->xb.data() + i0 * index->d);
}
}
}
bool successive_ids = index_flat != nullptr;
faiss::IndexShards *res =
new faiss::IndexShards(index->d, true,
successive_ids);
for (int i = 0; i < n; i++) {
res->add_shard(shards[i]);
}
res->own_fields = true;
FAISS_ASSERT(index->ntotal == res->ntotal);
return res;
}
Index *clone_Index(const Index *index) override {
long n = sub_cloners.size();
if (n == 1)
......@@ -236,54 +304,7 @@ struct ToGpuClonerMultiple: faiss::Cloner, GpuMultipleClonerOptions {
res->own_fields = true;
return res;
} else {
auto index_ivfpq =
dynamic_cast<const faiss::IndexIVFPQ *>(index);
auto index_ivfflat =
dynamic_cast<const faiss::IndexIVFFlat *>(index);
FAISS_THROW_IF_NOT_MSG (index_ivfpq || index_ivfflat,
"IndexShards implemented only for "
"IndexIVFFlat or IndexIVFPQ");
std::vector<faiss::Index*> shards(n);
for(long i = 0; i < n; i++) {
// make a shallow copy
if(reserveVecs)
sub_cloners[i].reserveVecs =
(reserveVecs + n - 1) / n;
if (index_ivfpq) {
faiss::IndexIVFPQ idx2(
index_ivfpq->quantizer, index_ivfpq->d,
index_ivfpq->nlist, index_ivfpq->code_size,
index_ivfpq->pq.nbits);
idx2.metric_type = index_ivfpq->metric_type;
idx2.pq = index_ivfpq->pq;
idx2.nprobe = index_ivfpq->nprobe;
idx2.use_precomputed_table = 0;
idx2.is_trained = index->is_trained;
copy_ivf_shard (index_ivfpq, &idx2, n, i);
shards[i] = sub_cloners[i].clone_Index(&idx2);
} else if (index_ivfflat) {
faiss::IndexIVFFlat idx2(
index_ivfflat->quantizer, index->d,
index_ivfflat->nlist, index_ivfflat->metric_type);
idx2.nprobe = index_ivfflat->nprobe;
idx2.nprobe = index_ivfflat->nprobe;
copy_ivf_shard (index_ivfflat, &idx2, n, i);
shards[i] = sub_cloners[i].clone_Index(&idx2);
}
}
faiss::IndexShards *res =
new faiss::IndexShards(index->d, true, false);
for (int i = 0; i < n; i++) {
res->add_shard(shards[i]);
}
res->own_fields = true;
FAISS_ASSERT(index->ntotal == res->ntotal);
return res;
return clone_Index_to_shards (index);
}
} else if(auto miq = dynamic_cast<const MultiIndexQuantizer *>(index)) {
if (verbose) {
......
gpu/StandardGpuResources.cpp
View file @ 76bec0b5
......@@ -30,7 +30,8 @@ StandardGpuResources::StandardGpuResources() :
tempMemFraction_(kDefaultTempMemFraction),
tempMemSize_(0),
useFraction_(true),
pinnedMemSize_(kDefaultPinnedMemoryAllocation) {
pinnedMemSize_(kDefaultPinnedMemoryAllocation),
cudaMallocWarning_(true) {
}
StandardGpuResources::~StandardGpuResources() {
......@@ -74,6 +75,7 @@ StandardGpuResources::~StandardGpuResources() {
void
StandardGpuResources::noTempMemory() {
setTempMemory(0);
setCudaMallocWarning(false);
}
void
......@@ -117,6 +119,15 @@ StandardGpuResources::setDefaultNullStreamAllDevices() {
}
}
void
StandardGpuResources::setCudaMallocWarning(bool b) {
cudaMallocWarning_ = b;
for (auto& v : memory_) {
v.second->setCudaMallocWarning(b);
}
}
void
StandardGpuResources::initializeForDevice(int device) {
// Use default streams as a marker for whether or not a certain
......@@ -195,9 +206,12 @@ StandardGpuResources::initializeForDevice(int device) {
}
FAISS_ASSERT(memory_.count(device) == 0);
memory_.emplace(device,
std::unique_ptr<StackDeviceMemory>(
new StackDeviceMemory(device, toAlloc)));
auto mem = std::unique_ptr<StackDeviceMemory>(
new StackDeviceMemory(device, toAlloc));
mem->setCudaMallocWarning(cudaMallocWarning_);
memory_.emplace(device, std::move(mem));
}
cublasHandle_t
......
gpu/StandardGpuResources.h
View file @ 76bec0b5
......@@ -48,6 +48,10 @@ class StandardGpuResources : public GpuResources {
/// for all devices
void setDefaultNullStreamAllDevices();
/// Enable or disable the warning about not having enough temporary memory
/// when cudaMalloc gets called
void setCudaMallocWarning(bool b);
public:
/// Internal system calls
void initializeForDevice(int device) override;
......@@ -100,6 +104,9 @@ class StandardGpuResources : public GpuResources {
/// Amount of pinned memory we should allocate
size_t pinnedMemSize_;
/// Whether or not a warning upon cudaMalloc is generated
bool cudaMallocWarning_;
};
} } // namespace
gpu/impl/IVFFlatScan.cu
View file @ 76bec0b5
This diff is collapsed.
gpu/test/TestGpuIndexFlat.cpp
View file @ 76bec0b5
......@@ -107,8 +107,6 @@ void testFlat(const TestFlatOptions& opt) {
TEST(TestGpuIndexFlat, IP_Float32) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
TestFlatOptions opt;
opt.useL2 = false;
opt.useFloat16 = false;
......@@ -123,8 +121,6 @@ TEST(TestGpuIndexFlat, IP_Float32) {
TEST(TestGpuIndexFlat, L2_Float32) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
TestFlatOptions opt;
opt.useL2 = true;
opt.useFloat16 = false;
......@@ -140,8 +136,6 @@ TEST(TestGpuIndexFlat, L2_Float32) {
// test specialized k == 1 codepath
TEST(TestGpuIndexFlat, L2_Float32_K1) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
TestFlatOptions opt;
opt.useL2 = true;
opt.useFloat16 = false;
......@@ -154,8 +148,6 @@ TEST(TestGpuIndexFlat, L2_Float32_K1) {
TEST(TestGpuIndexFlat, IP_Float16) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
TestFlatOptions opt;
opt.useL2 = false;
opt.useFloat16 = true;
......@@ -170,8 +162,6 @@ TEST(TestGpuIndexFlat, IP_Float16) {
TEST(TestGpuIndexFlat, L2_Float16) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
TestFlatOptions opt;
opt.useL2 = true;
opt.useFloat16 = true;
......@@ -187,8 +177,6 @@ TEST(TestGpuIndexFlat, L2_Float16) {
// test specialized k == 1 codepath
TEST(TestGpuIndexFlat, L2_Float16_K1) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
TestFlatOptions opt;
opt.useL2 = true;
opt.useFloat16 = true;
......@@ -201,15 +189,13 @@ TEST(TestGpuIndexFlat, L2_Float16_K1) {
// test tiling along a huge vector set
TEST(TestGpuIndexFlat, L2_Tiling) {
for (int tries = 0; tries < 3; ++tries) {
faiss::gpu::newTestSeed();
for (int tries = 0; tries < 2; ++tries) {
TestFlatOptions opt;
opt.useL2 = true;
opt.useFloat16 = false;
opt.useTransposed = false;
opt.numVecsOverride = 1000000;
opt.numQueriesOverride = 8;
opt.numQueriesOverride = 4;
testFlat(opt);
......@@ -251,8 +237,6 @@ TEST(TestGpuIndexFlat, QueryEmpty) {
}
TEST(TestGpuIndexFlat, CopyFrom) {
faiss::gpu::newTestSeed();
int numVecs = faiss::gpu::randVal(100, 200);
int dim = faiss::gpu::randVal(1, 1000);
......@@ -293,8 +277,6 @@ TEST(TestGpuIndexFlat, CopyFrom) {
}
TEST(TestGpuIndexFlat, CopyTo) {
faiss::gpu::newTestSeed();
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
......@@ -375,3 +357,12 @@ TEST(TestGpuIndexFlat, UnifiedMemory) {
0.1f,
0.015f);
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
// just run with a fixed test seed
faiss::gpu::setTestSeed(100);
return RUN_ALL_TESTS();
}
gpu/test/TestGpuIndexIVFFlat.cpp
View file @ 76bec0b5
......@@ -72,8 +72,6 @@ void queryTest(faiss::MetricType metricType,
bool useFloat16,
int dimOverride = -1) {
for (int tries = 0; tries < 3; ++tries) {
faiss::gpu::newTestSeed();
Options opt;
opt.dim = dimOverride != -1 ? dimOverride : opt.dim;
......@@ -116,7 +114,7 @@ void queryTest(faiss::MetricType metricType,
// FIXME: the fp16 bounds are
// useless when math (the accumulator) is
// in fp16. Figure out another way to test
compFloat16 ? 0.99f : 0.1f,
compFloat16 ? 0.70f : 0.1f,
compFloat16 ? 0.65f : 0.015f);
}
}
......@@ -125,8 +123,6 @@ void addTest(faiss::MetricType metricType,
bool useFloat16CoarseQuantizer,
bool useFloat16) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
......@@ -176,8 +172,6 @@ void addTest(faiss::MetricType metricType,
void copyToTest(bool useFloat16CoarseQuantizer,
bool useFloat16) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
......@@ -226,8 +220,6 @@ void copyToTest(bool useFloat16CoarseQuantizer,
void copyFromTest(bool useFloat16CoarseQuantizer,
bool useFloat16) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
......@@ -391,8 +383,6 @@ TEST(TestGpuIndexIVFFlat, Float32_32_CopyTo) {
}
TEST(TestGpuIndexIVFFlat, Float32_negative) {
faiss::gpu::newTestSeed();
Options opt;
auto trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
......@@ -457,8 +447,6 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
//
TEST(TestGpuIndexIVFFlat, QueryNaN) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
......@@ -505,8 +493,6 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
}
TEST(TestGpuIndexIVFFlat, AddNaN) {
faiss::gpu::newTestSeed();
Options opt;
faiss::gpu::StandardGpuResources res;
......@@ -612,3 +598,12 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
0.1f,
0.015f);
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
// just run with a fixed test seed
faiss::gpu::setTestSeed(100);
return RUN_ALL_TESTS();
}
gpu/test/TestGpuIndexIVFPQ.cpp
View file @ 76bec0b5
......@@ -94,7 +94,7 @@ struct Options {
}
float getPctMaxDiffN() const {
return useFloat16 ? 0.05f : 0.015f;
return useFloat16 ? 0.05f : 0.02f;
}
int numAdd;
......@@ -114,8 +114,6 @@ struct Options {
TEST(TestGpuIndexIVFPQ, Query) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
......@@ -150,8 +148,6 @@ TEST(TestGpuIndexIVFPQ, Query) {
TEST(TestGpuIndexIVFPQ, Add) {
for (int tries = 0; tries < 5; ++tries) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
......@@ -187,8 +183,6 @@ TEST(TestGpuIndexIVFPQ, Add) {
}
TEST(TestGpuIndexIVFPQ, CopyTo) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
......@@ -240,8 +234,6 @@ TEST(TestGpuIndexIVFPQ, CopyTo) {
}
TEST(TestGpuIndexIVFPQ, CopyFrom) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
......@@ -291,8 +283,6 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
}
TEST(TestGpuIndexIVFPQ, QueryNaN) {
faiss::gpu::newTestSeed();
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
......@@ -342,8 +332,6 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
}
TEST(TestGpuIndexIVFPQ, AddNaN) {
faiss::gpu::newTestSeed();
Options opt;
faiss::gpu::StandardGpuResources res;
......@@ -450,3 +438,12 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
0.1f,
0.015f);
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
// just run with a fixed test seed
faiss::gpu::setTestSeed(100);
return RUN_ALL_TESTS();
}
gpu/test/TestGpuSelect.cu
View file @ 76bec0b5
......@@ -19,11 +19,6 @@
#include <unordered_map>
#include <vector>
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
void testForSize(int rows, int cols, int k, bool dir, bool warp) {
std::vector<float> v = faiss::gpu::randVecs(rows, cols);
faiss::gpu::HostTensor<float, 2, true> hostVal({rows, cols});
......@@ -184,3 +179,12 @@ TEST(TestGpuSelect, testExactWarp) {
testForSize(rows, cols, cols, dir, true);
}
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
// just run with a fixed test seed
faiss::gpu::setTestSeed(100);
return RUN_ALL_TESTS();
}
gpu/test/test_gpu_index.py
View file @ 76bec0b5
......@@ -112,3 +112,94 @@ class EvalIVFPQAccuracy(unittest.TestCase):
res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
faiss.GpuParameterSpace().set_index_parameter(gpu_index, "nprobe", 3)
class ReferencedObject(unittest.TestCase):
d = 16
xb = np.random.rand(256, d).astype('float32')
nlist = 128
def test_proxy(self):
index = faiss.IndexProxy()
for i in range(3):
sub_index = faiss.IndexFlatL2(self.d)
sub_index.add(self.xb)
index.addIndex(sub_index)
assert index.d == self.d
index.search(self.xb, 10)
def test_resources(self):
# this used to crash!
index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0,
faiss.IndexFlatL2(self.d))
index.add(self.xb)
def test_flat(self):
index = faiss.GpuIndexFlat(faiss.StandardGpuResources(),
self.d, faiss.METRIC_L2)
index.add(self.xb)
def test_ivfflat(self):
index = faiss.GpuIndexIVFFlat(
faiss.StandardGpuResources(),
self.d, self.nlist, faiss.METRIC_L2)
index.train(self.xb)
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)
class TestShardedFlat(unittest.TestCase):
def test_sharded(self):
d = 32
nb = 1000
nq = 200
k = 10
rs = np.random.RandomState(123)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
index_cpu = faiss.IndexFlatL2(d)
assert faiss.get_num_gpus() > 1
co = faiss.GpuMultipleClonerOptions()
co.shard = True
index = faiss.index_cpu_to_all_gpus(index_cpu, co, ngpu=2)
index.add(xb)
D, I = index.search(xq, k)
index_cpu.add(xb)
D_ref, I_ref = index_cpu.search(xq, k)
assert np.all(I == I_ref)
del index
index2 = faiss.index_cpu_to_all_gpus(index_cpu, co, ngpu=2)
D2, I2 = index2.search(xq, k)
assert np.all(I2 == I_ref)
try:
index2.add(xb)
except RuntimeError:
pass
else:
assert False, "this call should fail!"
if __name__ == '__main__':
unittest.main()
gpu/utils/StackDeviceMemory.cpp
View file @ 76bec0b5
......@@ -25,7 +25,8 @@ StackDeviceMemory::Stack::Stack(int d, size_t sz)
head_(nullptr),
mallocCurrent_(0),
highWaterMemoryUsed_(0),
highWaterMalloc_(0) {
highWaterMalloc_(0),
cudaMallocWarning_(true) {
DeviceScope s(device_);
cudaError_t err = cudaMalloc(&start_, size_);
......@@ -41,7 +42,11 @@ StackDeviceMemory::Stack::Stack(int d, void* p, size_t sz, bool isOwner)
start_((char*) p),
end_(((char*) p) + sz),
size_(sz),
head_((char*) p) {
head_((char*) p),
mallocCurrent_(0),
highWaterMemoryUsed_(0),
highWaterMalloc_(0),
cudaMallocWarning_(true) {
}
StackDeviceMemory::Stack::~Stack() {
......@@ -59,15 +64,18 @@ StackDeviceMemory::Stack::getSizeAvailable() const {
}
char*
StackDeviceMemory::Stack::getAlloc(size_t size, cudaStream_t stream) {
StackDeviceMemory::Stack::getAlloc(size_t size,
cudaStream_t stream) {
if (size > (end_ - head_)) {
// Too large for our stack
DeviceScope s(device_);
// Print our requested size before we attempt the allocation
fprintf(stderr, "WARN: increase temp memory to avoid cudaMalloc, "
"or decrease query/add size (alloc %zu B, highwater %zu B)\n",
size, highWaterMalloc_);
if (cudaMallocWarning_) {
// Print our requested size before we attempt the allocation
fprintf(stderr, "WARN: increase temp memory to avoid cudaMalloc, "
"or decrease query/add size (alloc %zu B, highwater %zu B)\n",
size, highWaterMalloc_);
}
char* p = nullptr;
auto err = cudaMalloc(&p, size);
......@@ -190,6 +198,11 @@ StackDeviceMemory::StackDeviceMemory(int device,
StackDeviceMemory::~StackDeviceMemory() {
}
void
StackDeviceMemory::setCudaMallocWarning(bool b) {
stack_.cudaMallocWarning_ = b;
}
int
StackDeviceMemory::getDevice() const {
return device_;
......
gpu/utils/StackDeviceMemory.h
View file @ 76bec0b5
......@@ -29,6 +29,10 @@ class StackDeviceMemory : public DeviceMemory {
~StackDeviceMemory() override;
/// Enable or disable the warning about not having enough temporary memory
/// when cudaMalloc gets called
void setCudaMallocWarning(bool b);
int getDevice() const override;
DeviceMemoryReservation getMemory(cudaStream_t stream,
......@@ -111,6 +115,9 @@ class StackDeviceMemory : public DeviceMemory {
/// What's the high water mark in terms of memory allocated via
/// cudaMalloc?
size_t highWaterMalloc_;
/// Whether or not a warning upon cudaMalloc is generated
bool cudaMallocWarning_;
};
/// Our device
......
index_io.cpp
View file @ 76bec0b5
......@@ -226,7 +226,7 @@ static void write_ScalarQuantizer (
WRITEVECTOR (ivsc->trained);
}
static void write_InvertedLists (const InvertedLists *ils, IOWriter *f) {
void write_InvertedLists (const InvertedLists *ils, IOWriter *f) {
if (ils == nullptr) {
uint32_t h = fourcc ("il00");
WRITE1 (h);
......
index_io.h
View file @ 76bec0b5
......@@ -26,6 +26,7 @@ struct IndexIVF;
struct ProductQuantizer;
struct IOReader;
struct IOWriter;
struct InvertedLists;
void write_index (const Index *idx, const char *fname);
void write_index (const Index *idx, FILE *f);
......@@ -35,8 +36,7 @@ void write_index_binary (const IndexBinary *idx, const char *fname);
void write_index_binary (const IndexBinary *idx, FILE *f);
void write_index_binary (const IndexBinary *idx, IOWriter *writer);
const int IO_FLAG_MMAP = 1;
const int IO_FLAG_MMAP = 1; // try to memmap if possible
const int IO_FLAG_READ_ONLY = 2;
Index *read_index (const char *fname, int io_flags = 0);
......@@ -47,14 +47,14 @@ IndexBinary *read_index_binary (const char *fname, int io_flags = 0);
IndexBinary *read_index_binary (FILE * f, int io_flags = 0);
IndexBinary *read_index_binary (IOReader *reader, int io_flags = 0);
void write_VectorTransform (const VectorTransform *vt, const char *fname);
VectorTransform *read_VectorTransform (const char *fname);
ProductQuantizer * read_ProductQuantizer (const char*fname);
void write_ProductQuantizer (const ProductQuantizer*pq, const char *fname);
void write_InvertedLists (const InvertedLists *ils, IOWriter *f);
InvertedLists *read_InvertedLists (IOReader *reader, int io_flags = 0);
/* cloning functions */
Index *clone_index (const Index *);
......
python/faiss.py
View file @ 76bec0b5
......@@ -306,6 +306,88 @@ for symbol in dir(this_module):
handle_ParameterSpace(the_class)
###########################################
# Add Python references to objects
# we do this at the Python class wrapper level.
###########################################
def add_ref_in_constructor(the_class, parameter_no):
# adds a reference to parameter parameter_no in self
# so that that parameter does not get deallocated before self
original_init = the_class.__init__
def replacement_init(self, *args):
original_init(self, *args)
self.referenced_objects = [args[parameter_no]]
def replacement_init_multiple(self, *args):
original_init(self, *args)
pset = parameter_no[len(args)]
self.referenced_objects = [args[no] for no in pset]
if type(parameter_no) == dict:
# a list of parameters to keep, depending on the number of arguments
the_class.__init__ = replacement_init_multiple
else:
the_class.__init__ = replacement_init
def add_ref_in_method(the_class, method_name, parameter_no):
original_method = getattr(the_class, method_name)
def replacement_method(self, *args):
ref = args[parameter_no]
if not hasattr(self, 'referenced_objects'):
self.referenced_objects = [ref]
else:
self.referenced_objects.append(ref)
return original_method(self, *args)
setattr(the_class, method_name, replacement_method)
def add_ref_in_function(function_name, parameter_no):
# assumes the function returns an object
original_function = getattr(this_module, function_name)
def replacement_function(*args):
result = original_function(*args)
ref = args[parameter_no]
result.referenced_objects = [ref]
return result
setattr(this_module, function_name, replacement_function)
add_ref_in_constructor(IndexIVFFlat, 0)
add_ref_in_constructor(IndexIVFFlatDedup, 0)
add_ref_in_constructor(IndexPreTransform, {2: [0, 1], 1: [0]})
add_ref_in_method(IndexPreTransform, 'prepend_transform', 0)
add_ref_in_constructor(IndexIVFPQ, 0)
add_ref_in_constructor(IndexIVFPQR, 0)
add_ref_in_constructor(Index2Layer, 0)
add_ref_in_constructor(Level1Quantizer, 0)
add_ref_in_constructor(IndexIVFScalarQuantizer, 0)
add_ref_in_constructor(IndexIDMap, 0)
add_ref_in_constructor(IndexIDMap2, 0)
add_ref_in_method(IndexShards, 'add_shard', 0)
add_ref_in_constructor(IndexRefineFlat, 0)
add_ref_in_constructor(IndexBinaryIVF, 0)
add_ref_in_constructor(IndexBinaryFromFloat, 0)
if hasattr(this_module, 'IndexProxy'):
add_ref_in_method(IndexProxy, 'addIndex', 0)
# seems really marginal...
# remove_ref_from_method(IndexProxy, 'removeIndex', 0)
# handle all the GPUResources refs
add_ref_in_function('index_cpu_to_gpu', 0)
add_ref_in_constructor(GpuIndexFlat, 0)
add_ref_in_constructor(GpuIndexIVFFlat, 0)
add_ref_in_constructor(GpuIndexIVFPQ, 0)
###########################################
# GPU functions
###########################################
def index_cpu_to_gpu_multiple_py(resources, index, co=None):
"""builds the C++ vectors for the GPU indices and the
resources. Handles the common case where the resources are assigned to
......@@ -315,18 +397,22 @@ def index_cpu_to_gpu_multiple_py(resources, index, co=None):
for i, res in enumerate(resources):
vdev.push_back(i)
vres.push_back(res)
return index_cpu_to_gpu_multiple(vres, vdev, index, co)
index = index_cpu_to_gpu_multiple(vres, vdev, index, co)
index.referenced_objects = resources
return index
def index_cpu_to_all_gpus(index, co=None, ngpu=-1):
if ngpu == -1:
ngpu = get_num_gpus()
res = [StandardGpuResources() for i in range(ngpu)]
index2 = index_cpu_to_gpu_multiple_py(res, index, co)
index2.dont_dealloc = res
return index2
###########################################
# numpy array / std::vector conversions
###########################################
# mapping from vector names in swigfaiss.swig and the numpy dtype names
vector_name_map = {
'Float': 'float32',
......@@ -365,39 +451,9 @@ def copy_array_to_vector(a, v):
memcpy(v.data(), swig_ptr(a), a.nbytes)
class Kmeans:
def __init__(self, d, k, niter=25, verbose=False, spherical = False):
self.d = d
self.k = k
self.cp = ClusteringParameters()
self.cp.niter = niter
self.cp.verbose = verbose
self.cp.spherical = spherical
self.centroids = None
def train(self, x):
assert x.flags.contiguous
n, d = x.shape
assert d == self.d
clus = Clustering(d, self.k, self.cp)
if self.cp.spherical:
self.index = IndexFlatIP(d)
else:
self.index = IndexFlatL2(d)
clus.train(x, self.index)
centroids = vector_float_to_array(clus.centroids)
self.centroids = centroids.reshape(self.k, d)
self.obj = vector_float_to_array(clus.obj)
return self.obj[-1]
def assign(self, x):
assert self.centroids is not None, "should train before assigning"
index = IndexFlatL2(self.d)
index.add(self.centroids)
D, I = index.search(x, 1)
return D.ravel(), I.ravel()
###########################################
# Wrapper for a few functions
###########################################
def kmin(array, k):
"""return k smallest values (and their indices) of the lines of a
......@@ -480,3 +536,42 @@ def replacement_map_search_multiple(self, keys):
replace_method(MapLong2Long, 'add', replacement_map_add)
replace_method(MapLong2Long, 'search_multiple', replacement_map_search_multiple)
###########################################
# Kmeans object
###########################################
class Kmeans:
def __init__(self, d, k, niter=25, verbose=False, spherical = False):
self.d = d
self.k = k
self.cp = ClusteringParameters()
self.cp.niter = niter
self.cp.verbose = verbose
self.cp.spherical = spherical
self.centroids = None
def train(self, x):
assert x.flags.contiguous
n, d = x.shape
assert d == self.d
clus = Clustering(d, self.k, self.cp)
if self.cp.spherical:
self.index = IndexFlatIP(d)
else:
self.index = IndexFlatL2(d)
clus.train(x, self.index)
centroids = vector_float_to_array(clus.centroids)
self.centroids = centroids.reshape(self.k, d)
self.obj = vector_float_to_array(clus.obj)
return self.obj[-1]
def assign(self, x):
assert self.centroids is not None, "should train before assigning"
index = IndexFlatL2(self.d)
index.add(self.centroids)
D, I = index.search(x, 1)
return D.ravel(), I.ravel()
python/swigfaiss.py
View file @ 76bec0b5
This diff is collapsed.
python/swigfaiss.swig
View file @ 76bec0b5
......@@ -21,9 +21,11 @@
%module swigfaiss;
#endif
// fbode SWIG fails on warnings, so make them non fatal
#pragma SWIG nowarn=321
#pragma SWIG nowarn=403
#pragma SWIG nowarn=325
#pragma SWIG nowarn=389
typedef unsigned long uint64_t;
typedef uint64_t size_t;
......@@ -85,6 +87,7 @@ extern "C" {
#include "IndexBinaryFlat.h"
#include "IndexBinaryIVF.h"
#include "IndexBinaryFromFloat.h"
#include "index_io.h"
......@@ -155,7 +158,6 @@ namespace std {
%ignore faiss::hamming;
/*******************************************************************
* Parse headers
*******************************************************************/
......@@ -258,6 +260,7 @@ int get_num_gpus()
%include "IndexLSH.h"
%include "PolysemousTraining.h"
%include "IndexPQ.h"
%include "InvertedLists.h"
%include "IndexIVF.h"
%include "IndexScalarQuantizer.h"
%include "IndexHNSW.h"
......@@ -270,6 +273,7 @@ int get_num_gpus()
%include "IndexBinary.h"
%include "IndexBinaryFlat.h"
%include "IndexBinaryIVF.h"
%include "IndexBinaryFromFloat.h"
......@@ -481,6 +485,7 @@ struct AsyncIndexSearchC {
%typemap(out) faiss::IndexBinary * {
DOWNCAST ( IndexBinaryIVF )
DOWNCAST ( IndexBinaryFlat )
DOWNCAST ( IndexBinaryFromFloat )
// default for non-recognized classes
DOWNCAST ( IndexBinary )
if ($1 == NULL)
......
python/swigfaiss_gpu.py
View file @ 76bec0b5
This diff is collapsed.
python/swigfaiss_gpu_wrap.cpp
View file @ 76bec0b5
This diff is collapsed.
python/swigfaiss_wrap.cpp
View file @ 76bec0b5
This diff is collapsed.
tests/test_binary_io.py
View file @ 76bec0b5
......@@ -38,7 +38,7 @@ class TestBinaryFlat(unittest.TestCase):
index.add(self.xb)
D, I = index.search(self.xq, 3)
tmpnam = tempfile.NamedTemporaryFile().name
_, tmpnam = tempfile.mkstemp()
try:
faiss.write_index_binary(index, tmpnam)
......@@ -75,7 +75,7 @@ class TestBinaryIVF(unittest.TestCase):
index.add(self.xb)
D, I = index.search(self.xq, 3)
tmpnam = tempfile.NamedTemporaryFile().name
_, tmpnam = tempfile.mkstemp()
try:
faiss.write_index_binary(index, tmpnam)
......
tests/test_build_blocks.py
View file @ 76bec0b5
......@@ -207,11 +207,31 @@ class TestOrthognalReconstruct(unittest.TestCase):
x = rs.rand(30, 20).astype('float32')
xt = lt.apply_py(x)
try:
xtt = lt.reverse_transform(xt)
lt.reverse_transform(xt)
except Exception:
pass
else:
self.assertFalse('should do an exception')
class TestMAdd(unittest.TestCase):
def test_1(self):
# try with dimensions that are multiples of 16 or not
rs = np.random.RandomState(123)
swig_ptr = faiss.swig_ptr
for dim in 16, 32, 20, 25:
for repeat in 1, 2, 3, 4, 5:
a = rs.rand(dim).astype('float32')
b = rs.rand(dim).astype('float32')
c = np.zeros(dim, dtype='float32')
bf = rs.uniform(5.0) - 2.5
idx = faiss.fvec_madd_and_argmin(
dim, swig_ptr(a), bf, swig_ptr(b),
swig_ptr(c))
ref_c = a + b * bf
assert np.abs(c - ref_c).max() < 1e-5
assert idx == ref_c.argmin()
if __name__ == '__main__':
unittest.main()
tests/test_index.py
View file @ 76bec0b5
......@@ -30,7 +30,7 @@ def get_dataset_2(d, nb, nt, nq):
"""
d1 = 10 # intrinsic dimension (more or less)
n = nb + nt + nq
rs = np.random.RandomState(1234)
rs = np.random.RandomState(1338)
x = rs.normal(size=(n, d1))
x = np.dot(x, rs.rand(d1, d))
# now we have a d1-dim ellipsoid in d-dimensional space
......
tests/test_index_accuracy.py 0 → 100644
View file @ 76bec0b5
# 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.
#! /usr/bin/env python2
import numpy as np
import unittest
import faiss
# translation of test_knn.lua
def random_unitary(n, d, seed):
x = faiss.randn(n * d, seed).reshape(n, d)
faiss.normalize_L2(x)
return x
class Randu10k:
def __init__(self):
self.nb = 10000
self.nq = 1000
self.nt = 10000
self.d = 128
self.xb = random_unitary(self.nb, self.d, 1)
self.xt = random_unitary(self.nt, self.d, 2)
self.xq = random_unitary(self.nq, self.d, 3)
dotprods = np.dot(self.xq, self.xb.T)
self.gt = dotprods.argmax(1)
self.k = 100
def launch(self, name, index):
if not index.is_trained:
index.train(self.xt)
index.add(self.xb)
return index.search(self.xq, self.k)
def evalres(self, res):
D, I = res
e = {}
for rank in 1, 10, 100:
e[rank] = (I[:, :rank] == self.gt.reshape(-1, 1)).sum() / float(self.nq)
return e
ev = Randu10k()
d = ev.d
# Parameters inverted indexes
ncentroids = int(4 * np.sqrt(ev.nb))
kprobe = int(np.sqrt(ncentroids))
# Parameters for LSH
nbits = d
# Parameters for indexes involving PQ
M = d / 8 # for PQ: #subquantizers
nbits_per_index = 8 # for PQ
class IndexAccuracy(unittest.TestCase):
def test_IndexFlatIP(self):
q = faiss.IndexFlatIP(d) # Ask inner product
res = ev.launch('FLAT / IP', q)
e = ev.evalres(res)
assert e[1] == 1.0
def test_IndexFlatL2(self):
q = faiss.IndexFlatL2(d)
res = ev.launch('FLAT / L2', q)
e = ev.evalres(res)
assert e[1] == 1.0
def test_ivf_kmeans(self):
ivfk = faiss.IndexIVFFlat(faiss.IndexFlatL2(d), d, ncentroids)
ivfk.nprobe = kprobe
res = ev.launch('IVF K-means', ivfk)
e = ev.evalres(res)
# should give 0.260 0.260 0.260
assert e[1] > 0.2
def test_indexLSH(self):
q = faiss.IndexLSH(d, nbits)
res = ev.launch('FLAT / LSH Cosine', q)
e = ev.evalres(res)
# should give 0.070 0.250 0.580
assert e[10] > 0.2
def test_IndexLSH_32_48(self):
# CHECK: the difference between 32 and 48 does not make much sense
for nbits2 in 32, 48:
q = faiss.IndexLSH(d, nbits2)
res = ev.launch('LSH half size', q)
e = ev.evalres(res)
# should give 0.003 0.019 0.108
assert e[10] > 0.018, e
def test_IndexPQ(self):
q = faiss.IndexPQ(d, M, nbits_per_index)
res = ev.launch('FLAT / PQ L2', q)
e = ev.evalres(res)
# should give 0.070 0.230 0.260
assert e[10] > 0.2
# Approximate search module: PQ with inner product distance
def test_IndexPQ_ip(self):
q = faiss.IndexPQ(d, M, nbits_per_index, faiss.METRIC_INNER_PRODUCT)
res = ev.launch('FLAT / PQ IP', q)
e = ev.evalres(res)
# should give 0.070 0.230 0.260
#(same result as regular PQ on normalized distances)
assert e[10] > 0.2
def test_IndexIVFPQ(self):
ivfpq = faiss.IndexIVFPQ(faiss.IndexFlatL2(d), d, ncentroids, M, 8)
ivfpq.nprobe = kprobe
res = ev.launch('IVF PQ', ivfpq)
e = ev.evalres(res)
# should give 0.070 0.230 0.260
assert e[10] > 0.2
# TODO: translate evaluation of nested
# Approximate search: PQ with full vector refinement
def test_IndexPQ_refined(self):
q = faiss.IndexPQ(d, M, nbits_per_index)
res = ev.launch('PQ non-refined', q)
e = ev.evalres(res)
q.reset()
rq = faiss.IndexRefineFlat(q)
res = ev.launch('PQ refined', rq)
e2 = ev.evalres(res)
assert e2[10] >= e[10]
rq.k_factor = 4
res = ev.launch('PQ refined*4', rq)
e3 = ev.evalres(res)
assert e3[10] >= e2[10]
def test_polysemous(self):
index = faiss.IndexPQ(d, M, nbits_per_index)
index.do_polysemous_training = True
# reduce nb iterations to speed up training for the test
index.polysemous_training.n_iter = 50000
index.polysemous_training.n_redo = 1
res = ev.launch('normal PQ', index)
e_baseline = ev.evalres(res)
index.search_type = faiss.IndexPQ.ST_polysemous
index.polysemous_ht = int(M / 16. * 58)
stats = faiss.cvar.indexPQ_stats
stats.reset()
res = ev.launch('Polysemous ht=%d' % index.polysemous_ht,
index)
e_polysemous = ev.evalres(res)
print(e_baseline, e_polysemous, index.polysemous_ht)
print(stats.n_hamming_pass, stats.ncode)
# The randu dataset is difficult, so we are not too picky on
# the results. Here we assert that we have < 10 % loss when
# computing full PQ on fewer than 20% of the data.
assert stats.n_hamming_pass < stats.ncode / 5
# Test disabled because difference is 0.17 on aarch64
# TODO check why???
# assert e_polysemous[10] > e_baseline[10] - 0.1
def test_ScalarQuantizer(self):
quantizer = faiss.IndexFlatL2(d)
ivfpq = faiss.IndexIVFScalarQuantizer(
quantizer, d, ncentroids,
faiss.ScalarQuantizer.QT_8bit)
ivfpq.nprobe = kprobe
res = ev.launch('IVF SQ', ivfpq)
e = ev.evalres(res)
# should give 0.234 0.236 0.236
assert e[10] > 0.235
tests/test_index_binary.py
View file @ 76bec0b5
......@@ -124,10 +124,23 @@ class TestBinaryIVF(unittest.TestCase):
(self.xt, self.xb, self.xq) = make_binary_dataset(d, nb, nt, nq)
index = faiss.IndexBinaryFlat(d)
index.add(self.xb)
Dref, Iref = index.search(self.xq, 1)
Dref, Iref = index.search(self.xq, 10)
self.Dref = Dref
def test_ivf_flat(self):
def test_ivf_flat_exhaustive(self):
d = self.xq.shape[1] * 8
quantizer = faiss.IndexBinaryFlat(d)
index = faiss.IndexBinaryIVF(quantizer, d, 8)
index.cp.min_points_per_centroid = 5 # quiet warning
index.nprobe = 8
index.train(self.xt)
index.add(self.xb)
Divfflat, _ = index.search(self.xq, 10)
np.testing.assert_array_equal(self.Dref, Divfflat)
def test_ivf_flat2(self):
d = self.xq.shape[1] * 8
quantizer = faiss.IndexBinaryFlat(d)
......@@ -136,9 +149,9 @@ class TestBinaryIVF(unittest.TestCase):
index.nprobe = 4
index.train(self.xt)
index.add(self.xb)
Divfflat, _ = index.search(self.xq, 1)
Divfflat, _ = index.search(self.xq, 10)
self.assertGreaterEqual((self.Dref == Divfflat).sum(), 448)
self.assertEqual((self.Dref == Divfflat).sum(), 4122)
......
tests/test_index_binary_from_float.py 0 → 100644
View file @ 76bec0b5
# 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.
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
import unittest
import faiss
def make_binary_dataset(d, nb, nt, nq):
assert d % 8 == 0
rs = np.random.RandomState(123)
x = rs.randint(256, size=(nb + nq + nt, int(d / 8))).astype('uint8')
return x[:nt], x[nt:-nq], x[-nq:]
def binary_to_float(x):
n, d = x.shape
x8 = x.reshape(n * d, -1)
c8 = 2 * ((x8 >> np.arange(8)) & 1).astype('int8') - 1
return c8.astype('float32').reshape(n, d * 8)
class TestIndexBinaryFromFloat(unittest.TestCase):
"""Use a binary index backed by a float index"""
def test_index_from_float(self):
d = 256
nt = 0
nb = 1500
nq = 500
(xt, xb, xq) = make_binary_dataset(d, nb, nt, nq)
index_ref = faiss.IndexFlatL2(d)
index_ref.add(binary_to_float(xb))
index = faiss.IndexFlatL2(d)
index_bin = faiss.IndexBinaryFromFloat(index)
index_bin.add(xb)
D_ref, I_ref = index_ref.search(binary_to_float(xq), 10)
D, I = index_bin.search(xq, 10)
np.testing.assert_allclose((D_ref / 4.0).astype('int32'), D)
def test_wrapped_quantizer(self):
d = 256
nt = 150
nb = 1500
nq = 500
(xt, xb, xq) = make_binary_dataset(d, nb, nt, nq)
nlist = 16
quantizer_ref = faiss.IndexBinaryFlat(d)
index_ref = faiss.IndexBinaryIVF(quantizer_ref, d, nlist)
index_ref.train(xt)
index_ref.add(xb)
unwrapped_quantizer = faiss.IndexFlatL2(d)
quantizer = faiss.IndexBinaryFromFloat(unwrapped_quantizer)
index = faiss.IndexBinaryIVF(quantizer, d, nlist)
index.train(xt)
index.add(xb)
D_ref, I_ref = index_ref.search(xq, 10)
D, I = index.search(xq, 10)
np.testing.assert_array_equal(D_ref, D)
def test_wrapped_quantizer_IMI(self):
d = 256
nt = 3500
nb = 10000
nq = 500
(xt, xb, xq) = make_binary_dataset(d, nb, nt, nq)
index_ref = faiss.IndexBinaryFlat(d)
index_ref.add(xb)
nlist_exp = 6
nlist = 2 ** (2 * nlist_exp)
float_quantizer = faiss.MultiIndexQuantizer(d, 2, nlist_exp)
wrapped_quantizer = faiss.IndexBinaryFromFloat(float_quantizer)
wrapped_quantizer.train(xt)
assert nlist == float_quantizer.ntotal
index = faiss.IndexBinaryIVF(wrapped_quantizer, d,
float_quantizer.ntotal)
index.nprobe = 2048
assert index.is_trained
index.add(xb)
D_ref, I_ref = index_ref.search(xq, 10)
D, I = index.search(xq, 10)
recall = sum(gti[0] in Di[:10] for gti, Di in zip(D_ref, D)) \
/ float(D_ref.shape[0])
assert recall > 0.82, "recall = %g" % recall
def test_wrapped_quantizer_HNSW(self):
faiss.omp_set_num_threads(1)
def bin2float(v):
def byte2float(byte):
return np.array([-1.0 + 2.0 * (byte & (1 << b) != 0)
for b in range(0, 8)])
return np.hstack([byte2float(byte) for byte in v]).astype('float32')
def floatvec2nparray(v):
return np.array([np.float32(v.at(i)) for i in range(0, v.size())]) \
.reshape(-1, d)
d = 256
nt = 12800
nb = 10000
nq = 500
(xt, xb, xq) = make_binary_dataset(d, nb, nt, nq)
index_ref = faiss.IndexBinaryFlat(d)
index_ref.add(xb)
nlist = 256
clus = faiss.Clustering(d, nlist)
clus_index = faiss.IndexFlatL2(d)
xt_f = np.array([bin2float(v) for v in xt])
clus.train(xt_f, clus_index)
centroids = floatvec2nparray(clus.centroids)
hnsw_quantizer = faiss.IndexHNSWFlat(d, 32)
hnsw_quantizer.add(centroids)
hnsw_quantizer.is_trained = True
wrapped_quantizer = faiss.IndexBinaryFromFloat(hnsw_quantizer)
assert nlist == hnsw_quantizer.ntotal
assert nlist == wrapped_quantizer.ntotal
assert wrapped_quantizer.is_trained
index = faiss.IndexBinaryIVF(wrapped_quantizer, d,
hnsw_quantizer.ntotal)
index.nprobe = 128
assert index.is_trained
index.add(xb)
D_ref, I_ref = index_ref.search(xq, 10)
D, I = index.search(xq, 10)
recall = sum(gti[0] in Di[:10] for gti, Di in zip(D_ref, D)) \
/ float(D_ref.shape[0])
assert recall > 0.77, "recall = %g" % recall
tests/test_index_composite.py
View file @ 76bec0b5
......@@ -316,9 +316,7 @@ class TestRareIO(unittest.TestCase):
if in_pretransform:
# make sure it still works when wrapped in an IndexPreTransform
tmp = index1
index1 = faiss.IndexPreTransform(index1)
index1.dont_dealloc_me = tmp
index1.train(xt)
index1.add(xb)
......
tests/test_merge.cpp
View file @ 76bec0b5
......@@ -18,46 +18,13 @@
#include <faiss/FaissAssert.h>
#include <faiss/VectorTransform.h>
#include <faiss/OnDiskInvertedLists.h>
#include <faiss/IVFlib.h>
namespace {
// Main function to test
// Merge index1 into index0. Works on IndexIVF's and IndexIVF's
// embedded in a IndexPreTransform
void merge_into(faiss::Index *index0, faiss::Index *index1, bool shift_ids) {
FAISS_THROW_IF_NOT (index0->d == index1->d);
faiss::IndexIVF *ivf0 = dynamic_cast<faiss::IndexIVF *>(index0);
faiss::IndexIVF *ivf1 = dynamic_cast<faiss::IndexIVF *>(index1);
if (!ivf0) {
faiss::IndexPreTransform *pt0 = dynamic_cast<faiss::IndexPreTransform *>(index0);
faiss::IndexPreTransform *pt1 = dynamic_cast<faiss::IndexPreTransform *>(index1);
// minimal sanity check
FAISS_THROW_IF_NOT (pt0 && pt1);
FAISS_THROW_IF_NOT (pt0->chain.size() == pt1->chain.size());
for (int i = 0; i < pt0->chain.size(); i++) {
FAISS_THROW_IF_NOT (typeid(pt0->chain[i]) == typeid(pt1->chain[i]));
}
ivf0 = dynamic_cast<faiss::IndexIVF *>(pt0->index);
ivf1 = dynamic_cast<faiss::IndexIVF *>(pt1->index);
}
FAISS_THROW_IF_NOT (ivf0);
FAISS_THROW_IF_NOT (ivf1);
ivf0->merge_from (*ivf1, shift_ids ? ivf0->ntotal : 0);
// useful for IndexPreTransform
index0->ntotal = ivf0->ntotal;
index1->ntotal = ivf1->ntotal;
}
struct Tempfilename {
static pthread_mutex_t mutex;
......@@ -122,8 +89,6 @@ struct CommonData {
CommonData cd;
/// perform a search on shards, then merge and search again and
/// compare results.
int compare_merged (faiss::IndexShards *index_shards, bool shift_ids,
......@@ -142,7 +107,9 @@ int compare_merged (faiss::IndexShards *index_shards, bool shift_ids,
if (standard_merge) {
for (int i = 1; i < nindex; i++) {
merge_into(index_shards->at(0), index_shards->at(i), shift_ids);
faiss::ivflib::merge_into(
index_shards->at(0), index_shards->at(i),
shift_ids);
}
index_shards->sync_with_shard_indexes();
......@@ -275,7 +242,7 @@ TEST(MERGE, merge_flat_ondisk) {
EXPECT_EQ(ndiff, 0);
}
// non use ondisk specific merge
// now use ondisk specific merge
TEST(MERGE, merge_flat_ondisk_2) {
faiss::IndexShards index_shards(d, false, false);
index_shards.own_fields = true;
......
tests/test_pairs_decoding.cpp
View file @ 76bec0b5
......@@ -17,118 +17,17 @@
#include <faiss/IndexIVF.h>
#include <faiss/AutoTune.h>
#include <faiss/VectorTransform.h>
#include <faiss/IVFlib.h>
namespace {
typedef faiss::Index::idx_t idx_t;
/*************************************************************
* 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,
idx_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,
idx_t* labels,
idx_t* query_centroid_ids,
idx_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<idx_t> 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++) {
idx_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->invlists->get_single_id(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;
......@@ -162,7 +61,7 @@ std::unique_ptr<faiss::Index> make_index(const char *index_type,
* Test functions for a given index type
*************************************************************/
bool test_Search_centroid(const char *index_key) {
bool test_search_centroid(const char *index_key) {
std::vector<float> xb = make_data(nb); // database vectors
auto index = make_index(index_key, xb);
......@@ -171,9 +70,11 @@ bool test_Search_centroid(const char *index_key) {
the inverted list corresponding to its centroid */
std::vector<idx_t> centroid_ids (nb);
Search_centroid(index.get(), xb.data(), nb, centroid_ids.data());
faiss::ivflib::search_centroid(
index.get(), xb.data(), nb, centroid_ids.data());
const faiss::IndexIVF * ivf = get_IndexIVF(index.get());
const faiss::IndexIVF * ivf = faiss::ivflib::extract_index_ivf
(index.get());
for(int i = 0; i < nb; i++) {
bool found = false;
......@@ -197,9 +98,11 @@ int test_search_and_return_centroids(const char *index_key) {
auto index = make_index(index_key, xb);
std::vector<idx_t> centroid_ids (nb);
Search_centroid(index.get(), xb.data(), nb, centroid_ids.data());
faiss::ivflib::search_centroid(index.get(), xb.data(),
nb, centroid_ids.data());
faiss::IndexIVF * ivf = get_IndexIVF(index.get());
faiss::IndexIVF * ivf =
faiss::ivflib::extract_index_ivf (index.get());
ivf->nprobe = 4;
std::vector<float> xq = make_data(nq); // database vectors
......@@ -220,7 +123,7 @@ int test_search_and_return_centroids(const char *index_key) {
std::vector<idx_t> query_centroid_ids (nq);
std::vector<idx_t> result_centroid_ids (nq * k);
search_and_retrun_centroids(index.get(),
faiss::ivflib::search_and_return_centroids(index.get(),
nq, xq.data(), k,
newD.data(), newI.data(),
query_centroid_ids.data(),
......@@ -264,13 +167,13 @@ int test_search_and_return_centroids(const char *index_key) {
* Test entry points
*************************************************************/
TEST(test_Search_centroid, IVFFlat) {
bool ok = test_Search_centroid("IVF32,Flat");
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");
TEST(test_search_centroid, PCAIVFFlat) {
bool ok = test_search_centroid("PCA16,IVF32,Flat");
EXPECT_TRUE(ok);
}
......
tests/test_params_override.cpp
View file @ 76bec0b5
......@@ -17,6 +17,7 @@
#include <faiss/IndexIVF.h>
#include <faiss/IndexBinaryIVF.h>
#include <faiss/AutoTune.h>
#include <faiss/IVFlib.h>
using namespace faiss;
......@@ -45,19 +46,17 @@ std::vector<float> make_data(size_t n)
return database;
}
std::unique_ptr<IndexIVF> make_index(const char *index_type,
std::unique_ptr<Index> make_index(const char *index_type,
MetricType metric,
const std::vector<float> & x)
{
auto index = std::unique_ptr<IndexIVF>
(dynamic_cast<IndexIVF*>(index_factory(d, index_type, metric)));
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(IndexIVF *index, const float *xq) {
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);
......@@ -66,19 +65,12 @@ std::vector<idx_t> search_index(IndexIVF *index, const float *xq) {
}
std::vector<idx_t> search_index_with_params(
IndexIVF *index, const float *xq, IVFSearchParameters *params) {
Index *index, const float *xq, IVFSearchParameters *params) {
int k = 10;
std::vector<idx_t> I(k * nq);
std::vector<float> D(k * nq);
std::vector<idx_t> Iq(params->nprobe * nq);
std::vector<float> 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);
ivflib::search_with_parameters (index, nq, xq, k,
D.data(), I.data(), params);
return I;
}
......@@ -92,14 +84,15 @@ std::vector<idx_t> search_index_with_params(
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());
//index->train(nb, xb.data());
// index->add(nb, xb.data());
std::vector<float> xq = make_data(nq);
index->nprobe = 2;
ParameterSpace ps;
ps.set_index_parameter(index.get(), "nprobe", 2);
auto res2ref = search_index(index.get(), xq.data());
index->nprobe = 9;
ps.set_index_parameter(index.get(), "nprobe", 9);
auto res9ref = search_index(index.get(), xq.data());
index->nprobe = 1;
ps.set_index_parameter(index.get(), "nprobe", 1);
IVFSearchParameters params;
params.max_codes = 0;
......@@ -146,6 +139,13 @@ TEST(TPO, IVFSQ) {
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);
}
/*************************************************************
......
tests/test_referenced_objects.py 0 → 100644
View file @ 76bec0b5
# 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.
#! /usr/bin/env python2
"""make sure that the referenced objects are kept"""
import numpy as np
import unittest
import faiss
import sys
import gc
d = 10
xt = np.random.rand(100, d).astype('float32')
xb = np.random.rand(20, d).astype('float32')
class TestReferenced(unittest.TestCase):
def test_IndexIVF(self):
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, 10)
index.train(xt)
index.add(xb)
del quantizer
gc.collect()
index.add(xb)
def test_count_refs(self):
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, 10)
refc1 = sys.getrefcount(quantizer)
del index
gc.collect()
refc2 = sys.getrefcount(quantizer)
assert refc2 == refc1 - 1
def test_IndexIVF_2(self):
index = faiss.IndexIVFFlat(faiss.IndexFlatL2(d), d, 10)
index.train(xt)
index.add(xb)
def test_IndexPreTransform(self):
ltrans = faiss.NormalizationTransform(d)
sub_index = faiss.IndexFlatL2(d)
index = faiss.IndexPreTransform(ltrans, sub_index)
index.add(xb)
del ltrans
gc.collect()
index.add(xb)
del sub_index
gc.collect()
index.add(xb)
def test_IndexPreTransform_2(self):
sub_index = faiss.IndexFlatL2(d)
index = faiss.IndexPreTransform(sub_index)
ltrans = faiss.NormalizationTransform(d)
index.prepend_transform(ltrans)
index.add(xb)
del ltrans
gc.collect()
index.add(xb)
del sub_index
gc.collect()
index.add(xb)
def test_IDMap(self):
sub_index = faiss.IndexFlatL2(d)
index = faiss.IndexIDMap(sub_index)
index.add_with_ids(xb, np.arange(len(xb)))
del sub_index
gc.collect()
index.add_with_ids(xb, np.arange(len(xb)))
def test_shards(self):
index = faiss.IndexShards(d)
for i in range(3):
sub_index = faiss.IndexFlatL2(d)
sub_index.add(xb)
index.add_shard(sub_index)
gc.collect()
index.search(xb, 10)
if __name__ == '__main__':
unittest.main()
tests/test_sliding_ivf.cpp 0 → 100644
View file @ 76bec0b5
This diff is collapsed.
tests/test_transfer_invlists.cpp 0 → 100644
View file @ 76bec0b5
/**
* 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 <memory>
#include <cstdio>
#include <cstdlib>
#include <gtest/gtest.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/index_io.h>
#include <faiss/AuxIndexStructures.h>
#include <faiss/AutoTune.h>
#include <faiss/VectorTransform.h>
#include <faiss/utils.h>
#include <faiss/IVFlib.h>
using namespace faiss;
namespace {
// parameters to use for the test
int d = 64;
size_t nb = 1000;
size_t nq = 100;
size_t nt = 500;
int k = 10;
int nlist = 40;
typedef faiss::Index::idx_t idx_t;
std::vector<float> get_data (size_t nb, int seed) {
std::vector<float> x (nb * d);
float_randn (x.data(), nb * d, seed);
return x;
}
void test_index_type(const char *factory_string) {
// transfer inverted lists in nslice slices
int nslice = 3;
/****************************************************************
* trained reference index
****************************************************************/
std::unique_ptr<Index> trained (index_factory (d, factory_string));
{
auto xt = get_data (nt, 123);
trained->train (nt, xt.data());
}
// sample nq query vectors to check if results are the same
auto xq = get_data (nq, 818);
/****************************************************************
* source index
***************************************************************/
std::unique_ptr<Index> src_index (clone_index (trained.get()));
{ // add some data to source index
auto xb = get_data (nb, 245);
src_index->add (nb, xb.data());
}
ParameterSpace().set_index_parameter (src_index.get(), "nprobe", 4);
// remember reference search result on source index
std::vector<idx_t> Iref (nq * k);
std::vector<float> Dref (nq * k);
src_index->search (nq, xq.data(), k, Dref.data(), Iref.data());
/****************************************************************
* destination index -- should be replaced by source index
***************************************************************/
std::unique_ptr<Index> dst_index (clone_index (trained.get()));
{ // initial state: filled in with some garbage
int nb2 = nb + 10;
auto xb = get_data (nb2, 366);
dst_index->add (nb2, xb.data());
}
std::vector<idx_t> Inew (nq * k);
std::vector<float> Dnew (nq * k);
ParameterSpace().set_index_parameter (dst_index.get(), "nprobe", 4);
// transfer from source to destination in nslice slices
for (int sl = 0; sl < nslice; sl++) {
// so far, the indexes are different
dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
EXPECT_TRUE (Iref != Inew);
EXPECT_TRUE (Dref != Dnew);
// range of inverted list indices to transfer
long i0 = sl * nlist / nslice;
long i1 = (sl + 1) * nlist / nslice;
std::vector<uint8_t> data_to_transfer;
{
std::unique_ptr<ArrayInvertedLists> il
(ivflib::get_invlist_range (src_index.get(), i0, i1));
// serialize inverted lists
VectorIOWriter wr;
write_InvertedLists (il.get(), &wr);
data_to_transfer.swap (wr.data);
}
// transfer data here from source machine to dest machine
{
VectorIOReader reader;
reader.data.swap (data_to_transfer);
// deserialize inverted lists
std::unique_ptr<ArrayInvertedLists> il
(dynamic_cast<ArrayInvertedLists *>
(read_InvertedLists (&reader)));
// swap inverted lists. Block searches here!
{
ivflib::set_invlist_range (dst_index.get(), i0, i1, il.get());
}
}
}
EXPECT_EQ (dst_index->ntotal, src_index->ntotal);
// now, the indexes are the same
dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
EXPECT_TRUE (Iref == Inew);
EXPECT_TRUE (Dref == Dnew);
}
} // namespace
TEST(TRANS, IVFFlat) {
test_index_type ("IVF40,Flat");
}
TEST(TRANS, IVFFlatPreproc) {
test_index_type ("PCAR32,IVF40,Flat");
}
utils.cpp
View file @ 76bec0b5
This diff is collapsed.
utils.h
View file @ 76bec0b5
This diff is collapsed.
utils_simd.cpp 0 → 100644
View file @ 76bec0b5
This diff is collapsed.
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment