mirrors
/
faiss
mirror of https://github.com/facebookresearch/faiss.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
faiss/IndexBinaryIVF.cpp

627 lines
18 KiB
C++
Raw Blame History

/**
* 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.
*/
// Copyright 2004-present Facebook. All Rights Reserved
// -*- c++ -*-
#include "IndexBinaryIVF.h"
#include <cstdio>
#include <memory>
#include "hamming.h"
#include "utils.h"
#include "AuxIndexStructures.h"
#include "FaissAssert.h"
#include "IndexFlat.h"
namespace faiss {
IndexBinaryIVF::IndexBinaryIVF(IndexBinary *quantizer, size_t d, size_t nlist)
: IndexBinary(d),
invlists(new ArrayInvertedLists(nlist, code_size)),
own_invlists(true),
nprobe(1),
max_codes(0),
maintain_direct_map(false),
quantizer(quantizer),
nlist(nlist),
own_fields(false) {
FAISS_THROW_IF_NOT (d == quantizer->d);
is_trained = quantizer->is_trained && (quantizer->ntotal == nlist);
cp.niter = 10;
}
IndexBinaryIVF::IndexBinaryIVF()
: invlists(nullptr),
own_invlists(false),
nprobe(1),
max_codes(0),
maintain_direct_map(false),
quantizer(nullptr),
nlist(0),
own_fields(false) {}
void IndexBinaryIVF::add(idx_t n, const uint8_t *x) {
add_with_ids(n, x, nullptr);
}
void IndexBinaryIVF::add_with_ids(idx_t n, const uint8_t *x, const long *xids) {
add_core(n, x, xids, nullptr);
}
void IndexBinaryIVF::add_core(idx_t n, const uint8_t *x, const long *xids,
const long *precomputed_idx) {
FAISS_THROW_IF_NOT(is_trained);
assert(invlists);
FAISS_THROW_IF_NOT_MSG(!(maintain_direct_map && xids),
"cannot have direct map and add with ids");
const long * idx;
std::unique_ptr<long[]> scoped_idx;
if (precomputed_idx) {
idx = precomputed_idx;
} else {
scoped_idx.reset(new long[n]);
quantizer->assign(n, x, scoped_idx.get());
idx = scoped_idx.get();
}
long n_add = 0;
for (size_t i = 0; i < n; i++) {
long id = xids ? xids[i] : ntotal + i;
long list_no = idx[i];
if (list_no < 0)
continue;
const uint8_t *xi = x + i * code_size;
size_t offset = invlists->add_entry(list_no, id, xi);
if (maintain_direct_map)
direct_map.push_back(list_no << 32 | offset);
n_add++;
}
if (verbose) {
printf("IndexBinaryIVF::add_with_ids: added %ld / %ld vectors\n",
n_add, n);
}
ntotal += n_add;
}
void IndexBinaryIVF::make_direct_map(bool new_maintain_direct_map) {
// nothing to do
if (new_maintain_direct_map == maintain_direct_map)
return;
if (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);
for (long ofs = 0; ofs < list_size; ofs++) {
FAISS_THROW_IF_NOT_MSG(0 <= idlist[ofs] && idlist[ofs] < ntotal,
"direct map supported only for seuquential ids");
direct_map[idlist[ofs]] = key << 32 | ofs;
}
}
} else {
direct_map.clear();
}
maintain_direct_map = new_maintain_direct_map;
}
void IndexBinaryIVF::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels) const {
std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
invlists->prefetch_lists(idx.get(), n * nprobe);
search_preassigned(n, x, k, idx.get(), coarse_dis.get(),
distances, labels, false);
}
void IndexBinaryIVF::reconstruct(idx_t key, uint8_t *recons) const {
FAISS_THROW_IF_NOT_MSG(direct_map.size() == ntotal,
"direct map is not initialized");
long list_no = direct_map[key] >> 32;
long offset = direct_map[key] & 0xffffffff;
reconstruct_from_offset(list_no, offset, recons);
}
void IndexBinaryIVF::reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const {
FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
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);
for (long offset = 0; offset < list_size; offset++) {
long id = idlist[offset];
if (!(id >= i0 && id < i0 + ni)) {
continue;
}
uint8_t *reconstructed = recons + (id - i0) * d;
reconstruct_from_offset(list_no, offset, reconstructed);
}
}
}
void IndexBinaryIVF::search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
uint8_t *recons) const {
std::unique_ptr<idx_t[]> idx(new long[n * nprobe]);
std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
invlists->prefetch_lists(idx.get(), n * nprobe);
// search_preassigned() with `store_pairs` enabled to obtain the list_no
// and offset into `codes` for reconstruction
search_preassigned(n, x, k, idx.get(), coarse_dis.get(),
distances, labels, /* store_pairs */true);
for (idx_t i = 0; i < n; ++i) {
for (idx_t j = 0; j < k; ++j) {
idx_t ij = i * k + j;
idx_t key = labels[ij];
uint8_t *reconstructed = recons + ij * d;
if (key < 0) {
// Fill with NaNs
memset(reconstructed, -1, sizeof(*reconstructed) * d);
} else {
int list_no = key >> 32;
int offset = key & 0xffffffff;
// Update label to the actual id
labels[ij] = invlists->get_single_id(list_no, offset);
reconstruct_from_offset(list_no, offset, reconstructed);
}
}
}
}
void IndexBinaryIVF::reconstruct_from_offset(long list_no, long offset,
uint8_t *recons) const {
memcpy(recons, invlists->get_single_code(list_no, offset), code_size);
}
void IndexBinaryIVF::reset() {
direct_map.clear();
invlists->reset();
ntotal = 0;
}
long IndexBinaryIVF::remove_ids(const IDSelector& sel) {
FAISS_THROW_IF_NOT_MSG(!maintain_direct_map,
"direct map remove not implemented");
std::vector<long> toremove(nlist);
#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);
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));
} else {
j++;
}
}
toremove[i] = l0 - l;
}
// this will not run well in parallel on ondisk because of possible shrinks
long nremove = 0;
for (long i = 0; i < nlist; i++) {
if (toremove[i] > 0) {
nremove += toremove[i];
invlists->resize(
i, invlists->list_size(i) - toremove[i]);
}
}
ntotal -= nremove;
return nremove;
}
void IndexBinaryIVF::train(idx_t n, const uint8_t *x) {
if (verbose)
printf("Training level-1 quantizer\n");
train_q1(n, x, verbose);
is_trained = true;
}
double IndexBinaryIVF::imbalance_factor () const {
std::vector<int> hist(nlist);
for (int i = 0; i < nlist; i++) {
hist[i] = invlists->list_size(i);
}
return faiss::imbalance_factor(nlist, hist.data());
}
void IndexBinaryIVF::print_stats() const {
std::vector<int> sizes(40);
for (int i = 0; i < nlist; i++) {
for (int j = 0; j < sizes.size(); j++) {
if ((invlists->list_size(i) >> j) == 0) {
sizes[j]++;
break;
}
}
}
for (int i = 0; i < sizes.size(); i++) {
if (sizes[i]) {
printf("list size in < %d: %d instances\n", 1 << i, sizes[i]);
}
}
}
void IndexBinaryIVF::merge_from(IndexBinaryIVF &other, idx_t add_id) {
// 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;
}
invlists->add_entries(i, list_size, new_ids.data(),
oivf->get_codes(i));
}
oivf->resize(i, 0);
}
ntotal += other.ntotal;
other.ntotal = 0;
}
void IndexBinaryIVF::replace_invlists(InvertedLists *il, bool own) {
FAISS_THROW_IF_NOT(il->nlist == nlist &&
il->code_size == code_size);
if (own_invlists) {
delete invlists;
}
invlists = il;
own_invlists = 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)
printf("IVF quantizer does not need training.\n");
} else {
if (verbose)
printf("Training level-1 quantizer on %ld vectors in %dD\n", n, d);
Clustering clus(d, nlist, cp);
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);
}
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);
}
quantizer->add(clus.k, x_b.get());
quantizer->is_trained = true;
}
}
namespace {
template<class HammingComputer, bool store_pairs>
void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
size_t nx,
const uint8_t *x,
const long *keys,
int_maxheap_array_t *res,
const IVFSearchParameters *params) {
const size_t k = res->k;
size_t nlistv = 0, ndis = 0;
long nprobe = params ? params->nprobe : ivf.nprobe;
long max_codes = params ? params->max_codes : ivf.max_codes;
#pragma omp parallel for reduction(+: nlistv, ndis)
for (size_t i = 0; i < nx; i++) {
const uint8_t *xi = x + i * ivf.code_size;
const long * keysi = keys + i * nprobe;
int32_t * __restrict disi = res->get_val(i);
long * __restrict idxi = res->get_ids(i);
maxheap_heapify(k, disi, idxi);
size_t nscan = 0;
for (size_t ik = 0; ik < nprobe; ik++) {
long key = keysi[ik]; /* select the list */
if (key < 0) {
// not enough centroids for multiprobe
continue;
}
FAISS_THROW_IF_NOT_FMT (
key < (long) ivf.nlist,
"Invalid key=%ld at ik=%ld nlist=%ld\n",
key, ik, ivf.nlist);
nlistv++;
size_t list_size = ivf.invlists->list_size(key);
const uint8_t *list_vecs = (const uint8_t*)ivf.invlists->get_codes(key);
const Index::idx_t *ids = store_pairs
? nullptr
: ivf.invlists->get_ids(key);
HammingComputer hc(xi, ivf.code_size);
for (size_t j = 0; j < list_size; j++) {
const uint8_t * yj = list_vecs + ivf.code_size * j;
int32_t disij = hc.hamming(yj);
if (disij < disi[0]) {
maxheap_pop(k, disi, idxi);
long id = store_pairs ? (key << 32 | j) : ids[j];
maxheap_push(k, disi, idxi, disij, id);
}
}
nscan += list_size;
if (max_codes && nscan >= max_codes)
break;
}
ndis += nscan;
maxheap_reorder(k, disi, idxi);
}
indexIVF_stats.nq += nx;
indexIVF_stats.nlist += nlistv;
indexIVF_stats.ndis += ndis;
}
template<class HammingComputer, bool store_pairs>
void search_knn_hamming_count(const IndexBinaryIVF& ivf,
size_t nx,
const uint8_t *x,
const long *keys,
int k,
int32_t *distances,
long *labels,
const IVFSearchParameters *params) {
const int nBuckets = ivf.d + 1;
std::vector<int> all_counters(nx * nBuckets, 0);
std::unique_ptr<long[]> all_ids_per_dis(new long[nx * nBuckets * k]);
long nprobe = params ? params->nprobe : ivf.nprobe;
long max_codes = params ? params->max_codes : ivf.max_codes;
std::vector<HCounterState<HammingComputer>> cs;
for (size_t i = 0; i < nx; ++i) {
cs.push_back(HCounterState<HammingComputer>(
all_counters.data() + i * nBuckets,
all_ids_per_dis.get() + i * nBuckets * k,
x + i * ivf.code_size,
ivf.d,
k
));
}
size_t nlistv = 0, ndis = 0;
#pragma omp parallel for reduction(+: nlistv, ndis)
for (size_t i = 0; i < nx; i++) {
const long * keysi = keys + i * nprobe;
HCounterState<HammingComputer>& csi = cs[i];
size_t nscan = 0;
for (size_t ik = 0; ik < nprobe; ik++) {
long key = keysi[ik]; /* select the list */
if (key < 0) {
// not enough centroids for multiprobe
continue;
}
FAISS_THROW_IF_NOT_FMT (
key < (long) ivf.nlist,
"Invalid key=%ld at ik=%ld nlist=%ld\n",
key, ik, ivf.nlist);
nlistv++;
size_t list_size = ivf.invlists->list_size(key);
const uint8_t *list_vecs = (const uint8_t*)ivf.invlists->get_codes(key);
const Index::idx_t *ids = store_pairs
? nullptr
: ivf.invlists->get_ids(key);
for (size_t j = 0; j < list_size; j++) {
const uint8_t * yj = list_vecs + ivf.code_size * j;
long id = store_pairs ? (key << 32 | j) : ids[j];
csi.update_counter(yj, id);
}
nscan += list_size;
if (max_codes && nscan >= max_codes)
break;
}
ndis += nscan;
int nres = 0;
for (int b = 0; b < nBuckets && nres < k; b++) {
for (int l = 0; l < csi.counters[b] && nres < k; l++) {
labels[i * k + nres] = csi.ids_per_dis[b * k + l];
distances[i * k + nres] = b;
nres++;
}
}
}
indexIVF_stats.nq += nx;
indexIVF_stats.nlist += nlistv;
indexIVF_stats.ndis += ndis;
}
template<class HammingComputer>
void search_knn_hamming(const IndexBinaryIVF& ivf,
size_t nx,
const uint8_t *x,
const long *keys,
int k,
int32_t *distances,
long *labels,
bool store_pairs,
bool use_heap,
const IVFSearchParameters *params) {
if (use_heap) {
int_maxheap_array_t res = {
size_t(nx), size_t(k), labels, distances};
if (store_pairs) {
search_knn_hamming_heap<HammingComputer, true>(
ivf, nx, x, keys, &res, params);
} else {
search_knn_hamming_heap<HammingComputer, false>(
ivf, nx, x, keys, &res, params);
}
} else {
if (store_pairs) {
search_knn_hamming_count<HammingComputer, true>(
ivf, nx, x, keys, k, distances, labels, params);
} else {
search_knn_hamming_count<HammingComputer, false>(
ivf, nx, x, keys, k, distances, labels, params);
}
}
}
} // namespace
void IndexBinaryIVF::search_preassigned(idx_t n, const uint8_t *x, idx_t k,
const idx_t *idx,
const int32_t * /* coarse_dis */,
int32_t *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params
) const {
switch (code_size) {
#define HANNDLE_CS(cs) \
case cs: \
search_knn_hamming<HammingComputer ## cs>( \
*this, n, x, idx, k, distances, \
labels, store_pairs, use_heap, params); \
break;
HANNDLE_CS(4);
HANNDLE_CS(8);
HANNDLE_CS(16);
HANNDLE_CS(20);
HANNDLE_CS(32);
HANNDLE_CS(64);
#undef HANNDLE_CS
default:
if (code_size % 8 == 0) {
search_knn_hamming<HammingComputerM8>(
*this, n, x, idx, k, distances,
labels, store_pairs, use_heap, params);
} else if (code_size % 4 == 0) {
search_knn_hamming<HammingComputerM4>(
*this, n, x, idx, k, distances,
labels, store_pairs, use_heap, params);
} else {
search_knn_hamming<HammingComputerDefault>(
*this, n, x, idx, k, distances,
labels, store_pairs, use_heap, params);
}
break;
}
}
IndexBinaryIVF::~IndexBinaryIVF() {
if (own_invlists) {
delete invlists;
}
if (own_fields) {
delete quantizer;
}
}
} // namespace faiss
Reference in New Issue View Git Blame Copy Permalink