mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Facebook sync (May 2019) + relicense (#838)
Changelog: - changed license: BSD+Patents -> MIT - propagates exceptions raised in sub-indexes of IndexShards and IndexReplicas - support for searching several inverted lists in parallel (parallel_mode != 0) - better support for PQ codes where nbit != 8 or 16 - IVFSpectralHash implementation: spectral hash codes inside an IVF - 6-bit per component scalar quantizer (4 and 8 bit were already supported) - combinations of inverted lists: HStackInvertedLists and VStackInvertedLists - configurable number of threads for OnDiskInvertedLists prefetching (including 0=no prefetch) - more test and demo code compatible with Python 3 (print with parentheses) - refactored benchmark code: data loading is now in a single file
This commit is contained in:
Lucas Hosseini
+219
-232
@@ -1,8 +1,7 @@
|
||||
/**
|
||||
* Copyright (c) 2015-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
* Copyright (c) Facebook, Inc. and its affiliates.
|
||||
*
|
||||
* This source code is licensed under the BSD+Patents license found in the
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
@@ -43,288 +42,276 @@ void translate_labels (long n, idx_t *labels, long translation)
|
||||
*/
|
||||
|
||||
template <class IndexClass, class C>
|
||||
void merge_tables (long n, long k, long nshard,
|
||||
typename IndexClass::distance_t *distances,
|
||||
idx_t *labels,
|
||||
const typename IndexClass::distance_t *all_distances,
|
||||
idx_t *all_labels,
|
||||
const long *translations)
|
||||
{
|
||||
if(k == 0) {
|
||||
return;
|
||||
}
|
||||
using distance_t = typename IndexClass::distance_t;
|
||||
void
|
||||
merge_tables(long n, long k, long nshard,
|
||||
typename IndexClass::distance_t *distances,
|
||||
idx_t *labels,
|
||||
const std::vector<typename IndexClass::distance_t>& all_distances,
|
||||
const std::vector<idx_t>& all_labels,
|
||||
const std::vector<long>& translations) {
|
||||
if (k == 0) {
|
||||
return;
|
||||
}
|
||||
using distance_t = typename IndexClass::distance_t;
|
||||
|
||||
long stride = n * k;
|
||||
long stride = n * k;
|
||||
#pragma omp parallel
|
||||
{
|
||||
std::vector<int> buf (2 * nshard);
|
||||
int * pointer = buf.data();
|
||||
int * shard_ids = pointer + nshard;
|
||||
std::vector<distance_t> buf2 (nshard);
|
||||
distance_t * heap_vals = buf2.data();
|
||||
{
|
||||
std::vector<int> buf (2 * nshard);
|
||||
int * pointer = buf.data();
|
||||
int * shard_ids = pointer + nshard;
|
||||
std::vector<distance_t> buf2 (nshard);
|
||||
distance_t * heap_vals = buf2.data();
|
||||
#pragma omp for
|
||||
for (long i = 0; i < n; i++) {
|
||||
// the heap maps values to the shard where they are
|
||||
// produced.
|
||||
const distance_t *D_in = all_distances + i * k;
|
||||
const idx_t *I_in = all_labels + i * k;
|
||||
int heap_size = 0;
|
||||
for (long i = 0; i < n; i++) {
|
||||
// the heap maps values to the shard where they are
|
||||
// produced.
|
||||
const distance_t *D_in = all_distances.data() + i * k;
|
||||
const idx_t *I_in = all_labels.data() + i * k;
|
||||
int heap_size = 0;
|
||||
|
||||
for (long s = 0; s < nshard; s++) {
|
||||
pointer[s] = 0;
|
||||
if (I_in[stride * s] >= 0)
|
||||
heap_push<C> (++heap_size, heap_vals, shard_ids,
|
||||
D_in[stride * s], s);
|
||||
}
|
||||
|
||||
distance_t *D = distances + i * k;
|
||||
idx_t *I = labels + i * k;
|
||||
|
||||
for (int j = 0; j < k; j++) {
|
||||
if (heap_size == 0) {
|
||||
I[j] = -1;
|
||||
D[j] = C::neutral();
|
||||
} else {
|
||||
// pop best element
|
||||
int s = shard_ids[0];
|
||||
int & p = pointer[s];
|
||||
D[j] = heap_vals[0];
|
||||
I[j] = I_in[stride * s + p] + translations[s];
|
||||
|
||||
heap_pop<C> (heap_size--, heap_vals, shard_ids);
|
||||
p++;
|
||||
if (p < k && I_in[stride * s + p] >= 0)
|
||||
heap_push<C> (++heap_size, heap_vals, shard_ids,
|
||||
D_in[stride * s + p], s);
|
||||
}
|
||||
}
|
||||
for (long s = 0; s < nshard; s++) {
|
||||
pointer[s] = 0;
|
||||
if (I_in[stride * s] >= 0) {
|
||||
heap_push<C> (++heap_size, heap_vals, shard_ids,
|
||||
D_in[stride * s], s);
|
||||
}
|
||||
}
|
||||
|
||||
distance_t *D = distances + i * k;
|
||||
idx_t *I = labels + i * k;
|
||||
|
||||
for (int j = 0; j < k; j++) {
|
||||
if (heap_size == 0) {
|
||||
I[j] = -1;
|
||||
D[j] = C::neutral();
|
||||
} else {
|
||||
// pop best element
|
||||
int s = shard_ids[0];
|
||||
int & p = pointer[s];
|
||||
D[j] = heap_vals[0];
|
||||
I[j] = I_in[stride * s + p] + translations[s];
|
||||
|
||||
heap_pop<C> (heap_size--, heap_vals, shard_ids);
|
||||
p++;
|
||||
if (p < k && I_in[stride * s + p] >= 0) {
|
||||
heap_push<C> (++heap_size, heap_vals, shard_ids,
|
||||
D_in[stride * s + p], s);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<class IndexClass>
|
||||
void runOnIndexes(bool threaded,
|
||||
std::function<void(int no, IndexClass*)> f,
|
||||
std::vector<IndexClass *> indexes)
|
||||
{
|
||||
FAISS_THROW_IF_NOT_MSG(!indexes.empty(), "no shards in index");
|
||||
|
||||
if (!threaded) {
|
||||
for (int no = 0; no < indexes.size(); no++) {
|
||||
IndexClass *index = indexes[no];
|
||||
f(no, index);
|
||||
}
|
||||
} else {
|
||||
std::vector<std::unique_ptr<WorkerThread> > threads;
|
||||
std::vector<std::future<bool>> v;
|
||||
|
||||
for (int no = 0; no < indexes.size(); no++) {
|
||||
IndexClass *index = indexes[no];
|
||||
threads.emplace_back(new WorkerThread());
|
||||
WorkerThread *wt = threads.back().get();
|
||||
v.emplace_back(wt->add([no, index, f](){ f(no, index); }));
|
||||
}
|
||||
|
||||
// Blocking wait for completion
|
||||
for (auto& func : v) {
|
||||
func.get();
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
|
||||
template<class IndexClass>
|
||||
IndexShardsTemplate<IndexClass>::IndexShardsTemplate (idx_t d, bool threaded, bool successive_ids):
|
||||
IndexClass (d), own_fields (false),
|
||||
threaded (threaded), successive_ids (successive_ids)
|
||||
{
|
||||
|
||||
|
||||
template <typename IndexT>
|
||||
IndexShardsTemplate<IndexT>::IndexShardsTemplate(idx_t d,
|
||||
bool threaded,
|
||||
bool successive_ids)
|
||||
: ThreadedIndex<IndexT>(d, threaded),
|
||||
successive_ids(successive_ids) {
|
||||
}
|
||||
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::add_shard (IndexClass *idx)
|
||||
{
|
||||
shard_indexes.push_back (idx);
|
||||
sync_with_shard_indexes ();
|
||||
template <typename IndexT>
|
||||
IndexShardsTemplate<IndexT>::IndexShardsTemplate(int d,
|
||||
bool threaded,
|
||||
bool successive_ids)
|
||||
: ThreadedIndex<IndexT>(d, threaded),
|
||||
successive_ids(successive_ids) {
|
||||
}
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::sync_with_shard_indexes ()
|
||||
{
|
||||
if (shard_indexes.empty()) return;
|
||||
IndexClass * index0 = shard_indexes[0];
|
||||
this->d = index0->d;
|
||||
this->metric_type = index0->metric_type;
|
||||
this->is_trained = index0->is_trained;
|
||||
this->ntotal = index0->ntotal;
|
||||
for (int i = 1; i < shard_indexes.size(); i++) {
|
||||
IndexClass * index = shard_indexes[i];
|
||||
FAISS_THROW_IF_NOT (this->metric_type == index->metric_type);
|
||||
FAISS_THROW_IF_NOT (this->d == index->d);
|
||||
this->ntotal += index->ntotal;
|
||||
}
|
||||
template <typename IndexT>
|
||||
IndexShardsTemplate<IndexT>::IndexShardsTemplate(bool threaded,
|
||||
bool successive_ids)
|
||||
: ThreadedIndex<IndexT>(threaded),
|
||||
successive_ids(successive_ids) {
|
||||
}
|
||||
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::onAfterAddIndex(IndexT* index /* unused */) {
|
||||
sync_with_shard_indexes();
|
||||
}
|
||||
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::onAfterRemoveIndex(IndexT* index /* unused */) {
|
||||
sync_with_shard_indexes();
|
||||
}
|
||||
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::sync_with_shard_indexes() {
|
||||
if (!this->count()) {
|
||||
this->is_trained = false;
|
||||
this->ntotal = 0;
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
auto firstIndex = this->at(0);
|
||||
this->metric_type = firstIndex->metric_type;
|
||||
this->is_trained = firstIndex->is_trained;
|
||||
this->ntotal = firstIndex->ntotal;
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::train (idx_t n, const component_t *x)
|
||||
{
|
||||
auto train_func = [n, x](int no, IndexClass *index)
|
||||
{
|
||||
if (index->verbose)
|
||||
printf ("begin train shard %d on %ld points\n", no, n);
|
||||
index->train(n, x);
|
||||
if (index->verbose)
|
||||
printf ("end train shard %d\n", no);
|
||||
for (int i = 1; i < this->count(); ++i) {
|
||||
auto index = this->at(i);
|
||||
FAISS_THROW_IF_NOT(this->metric_type == index->metric_type);
|
||||
FAISS_THROW_IF_NOT(this->d == index->d);
|
||||
|
||||
this->ntotal += index->ntotal;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::train(idx_t n,
|
||||
const component_t *x) {
|
||||
auto fn =
|
||||
[n, x](int no, IndexT *index) {
|
||||
if (index->verbose) {
|
||||
printf("begin train shard %d on %ld points\n", no, n);
|
||||
}
|
||||
|
||||
index->train(n, x);
|
||||
|
||||
if (index->verbose) {
|
||||
printf("end train shard %d\n", no);
|
||||
}
|
||||
};
|
||||
|
||||
runOnIndexes<IndexClass> (threaded, train_func, shard_indexes);
|
||||
sync_with_shard_indexes ();
|
||||
this->runOnIndex(fn);
|
||||
sync_with_shard_indexes();
|
||||
}
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::add (idx_t n, const component_t *x)
|
||||
{
|
||||
add_with_ids (n, x, nullptr);
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::add(idx_t n,
|
||||
const component_t *x) {
|
||||
add_with_ids(n, x, nullptr);
|
||||
}
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::add_with_ids (idx_t n, const component_t * x, const idx_t *xids)
|
||||
{
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::add_with_ids(idx_t n,
|
||||
const component_t * x,
|
||||
const idx_t *xids) {
|
||||
|
||||
FAISS_THROW_IF_NOT_MSG(!(successive_ids && xids),
|
||||
"It makes no sense to pass in ids and "
|
||||
"request them to be shifted");
|
||||
FAISS_THROW_IF_NOT_MSG(!(successive_ids && xids),
|
||||
"It makes no sense to pass in ids and "
|
||||
"request them to be shifted");
|
||||
|
||||
if (successive_ids) {
|
||||
FAISS_THROW_IF_NOT_MSG(!xids,
|
||||
"It makes no sense to pass in ids and "
|
||||
"request them to be shifted");
|
||||
FAISS_THROW_IF_NOT_MSG(this->ntotal == 0,
|
||||
"when adding to IndexShards with sucessive_ids, "
|
||||
"only add() in a single pass is supported");
|
||||
if (successive_ids) {
|
||||
FAISS_THROW_IF_NOT_MSG(!xids,
|
||||
"It makes no sense to pass in ids and "
|
||||
"request them to be shifted");
|
||||
FAISS_THROW_IF_NOT_MSG(this->ntotal == 0,
|
||||
"when adding to IndexShards with sucessive_ids, "
|
||||
"only add() in a single pass is supported");
|
||||
}
|
||||
|
||||
idx_t nshard = this->count();
|
||||
const idx_t *ids = xids;
|
||||
|
||||
std::vector<idx_t> aids;
|
||||
|
||||
if (!ids && !successive_ids) {
|
||||
aids.resize(n);
|
||||
|
||||
for (idx_t i = 0; i < n; i++) {
|
||||
aids[i] = this->ntotal + i;
|
||||
}
|
||||
|
||||
long nshard = shard_indexes.size();
|
||||
const idx_t *ids = xids;
|
||||
ScopeDeleter<idx_t> del;
|
||||
if (!ids && !successive_ids) {
|
||||
idx_t *aids = new idx_t[n];
|
||||
for (idx_t i = 0; i < n; i++)
|
||||
aids[i] = this->ntotal + i;
|
||||
ids = aids;
|
||||
del.set (ids);
|
||||
}
|
||||
ids = aids.data();
|
||||
}
|
||||
|
||||
size_t components_per_vec =
|
||||
sizeof(component_t) == 1 ? (this->d + 7) / 8 : this->d;
|
||||
size_t components_per_vec =
|
||||
sizeof(component_t) == 1 ? (this->d + 7) / 8 : this->d;
|
||||
|
||||
auto add_func = [n, ids, x, nshard, components_per_vec]
|
||||
(int no, IndexClass *index) {
|
||||
auto fn =
|
||||
[n, ids, x, nshard, components_per_vec](int no, IndexT *index) {
|
||||
idx_t i0 = (idx_t) no * n / nshard;
|
||||
idx_t i1 = ((idx_t) no + 1) * n / nshard;
|
||||
auto x0 = x + i0 * components_per_vec;
|
||||
|
||||
idx_t i0 = no * n / nshard;
|
||||
idx_t i1 = (no + 1) * n / nshard;
|
||||
if (index->verbose) {
|
||||
printf ("begin add shard %d on %ld points\n", no, n);
|
||||
}
|
||||
|
||||
auto x0 = x + i0 * components_per_vec;
|
||||
if (ids) {
|
||||
index->add_with_ids (i1 - i0, x0, ids + i0);
|
||||
} else {
|
||||
index->add (i1 - i0, x0);
|
||||
}
|
||||
|
||||
if (index->verbose) {
|
||||
printf ("begin add shard %d on %ld points\n", no, n);
|
||||
}
|
||||
if (ids) {
|
||||
index->add_with_ids (i1 - i0, x0, ids + i0);
|
||||
} else {
|
||||
index->add (i1 - i0, x0);
|
||||
}
|
||||
if (index->verbose) {
|
||||
printf ("end add shard %d on %ld points\n", no, i1 - i0);
|
||||
}
|
||||
if (index->verbose) {
|
||||
printf ("end add shard %d on %ld points\n", no, i1 - i0);
|
||||
}
|
||||
};
|
||||
|
||||
runOnIndexes<IndexClass> (threaded, add_func, shard_indexes);
|
||||
this->runOnIndex(fn);
|
||||
|
||||
this->ntotal += n;
|
||||
// This is safe to do here because the current thread controls execution in
|
||||
// all threads, and nothing else is happening
|
||||
this->ntotal += n;
|
||||
}
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::reset ()
|
||||
{
|
||||
for (int i = 0; i < shard_indexes.size(); i++) {
|
||||
shard_indexes[i]->reset ();
|
||||
}
|
||||
sync_with_shard_indexes ();
|
||||
}
|
||||
template <typename IndexT>
|
||||
void
|
||||
IndexShardsTemplate<IndexT>::search(idx_t n,
|
||||
const component_t *x,
|
||||
idx_t k,
|
||||
distance_t *distances,
|
||||
idx_t *labels) const {
|
||||
long nshard = this->count();
|
||||
|
||||
template<class IndexClass>
|
||||
void IndexShardsTemplate<IndexClass>::search (
|
||||
idx_t n, const component_t *x, idx_t k,
|
||||
distance_t *distances, idx_t *labels) const
|
||||
{
|
||||
long nshard = shard_indexes.size();
|
||||
distance_t *all_distances = new distance_t [nshard * k * n];
|
||||
idx_t *all_labels = new idx_t [nshard * k * n];
|
||||
ScopeDeleter<distance_t> del (all_distances);
|
||||
ScopeDeleter<idx_t> del2 (all_labels);
|
||||
std::vector<distance_t> all_distances(nshard * k * n);
|
||||
std::vector<idx_t> all_labels(nshard * k * n);
|
||||
|
||||
auto query_func = [n, k, x, all_distances, all_labels]
|
||||
(int no, IndexClass *index) {
|
||||
auto fn =
|
||||
[n, k, x, &all_distances, &all_labels](int no, const IndexT *index) {
|
||||
if (index->verbose) {
|
||||
printf ("begin query shard %d on %ld points\n", no, n);
|
||||
}
|
||||
|
||||
if (index->verbose) {
|
||||
printf ("begin query shard %d on %ld points\n", no, n);
|
||||
}
|
||||
index->search (n, x, k,
|
||||
all_distances + no * k * n,
|
||||
all_labels + no * k * n);
|
||||
if (index->verbose) {
|
||||
printf ("end query shard %d\n", no);
|
||||
}
|
||||
index->search (n, x, k,
|
||||
all_distances.data() + no * k * n,
|
||||
all_labels.data() + no * k * n);
|
||||
|
||||
if (index->verbose) {
|
||||
printf ("end query shard %d\n", no);
|
||||
}
|
||||
};
|
||||
|
||||
runOnIndexes<IndexClass> (threaded, query_func, shard_indexes);
|
||||
this->runOnIndex(fn);
|
||||
|
||||
std::vector<long> translations (nshard, 0);
|
||||
if (successive_ids) {
|
||||
translations[0] = 0;
|
||||
for (int s = 0; s + 1 < nshard; s++)
|
||||
translations [s + 1] = translations [s] +
|
||||
shard_indexes [s]->ntotal;
|
||||
std::vector<long> translations(nshard, 0);
|
||||
|
||||
// Because we just called runOnIndex above, it is safe to access the sub-index
|
||||
// ntotal here
|
||||
if (successive_ids) {
|
||||
translations[0] = 0;
|
||||
|
||||
for (int s = 0; s + 1 < nshard; s++) {
|
||||
translations[s + 1] = translations[s] + this->at(s)->ntotal;
|
||||
}
|
||||
}
|
||||
|
||||
if (this->metric_type == METRIC_L2) {
|
||||
merge_tables<IndexClass, CMin<distance_t, int> > (
|
||||
n, k, nshard, distances, labels,
|
||||
all_distances, all_labels, translations.data ());
|
||||
} else {
|
||||
merge_tables<IndexClass, CMax<distance_t, int> > (
|
||||
n, k, nshard, distances, labels,
|
||||
all_distances, all_labels, translations.data ());
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
template<class IndexClass>
|
||||
IndexShardsTemplate<IndexClass>::~IndexShardsTemplate ()
|
||||
{
|
||||
if (own_fields) {
|
||||
for (int s = 0; s < shard_indexes.size(); s++)
|
||||
delete shard_indexes [s];
|
||||
}
|
||||
if (this->metric_type == METRIC_L2) {
|
||||
merge_tables<IndexT, CMin<distance_t, int>>(
|
||||
n, k, nshard, distances, labels,
|
||||
all_distances, all_labels, translations);
|
||||
} else {
|
||||
merge_tables<IndexT, CMax<distance_t, int>>(
|
||||
n, k, nshard, distances, labels,
|
||||
all_distances, all_labels, translations);
|
||||
}
|
||||
}
|
||||
|
||||
// explicit instanciations
|
||||
template struct IndexShardsTemplate<Index>;
|
||||
template struct IndexShardsTemplate<IndexBinary>;
|
||||
|
||||
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
Reference in New Issue
Block a user