facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

IndexFlatCodes: a single parent for all flat codecs (#2132)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2132

This diff adds the class IndexFlatCodes that becomes the parent of all "flat" encodings.
IndexPQ
IndexFlat
IndexAdditiveQuantizer
IndexScalarQuantizer
IndexLSH
Index2Layer

The other changes are:
- for IndexFlat, there is no vector<float> with the data anymore. It is replaced with a `get_xb()` function. This broke quite a few external codes, that this diff also attempts to fix.
- I/O functions needed to be adapted. This is done without changing the I/O format for any index.
- added a small contrib function to get the data from the IndexFlat
- the functionality has been made uniform, for example remove_ids and add are now in the parent class.

Eventually, we may support generic storage for flat indexes, similar to `InvertedLists`, eg to memmap the data, but this will again require a big change.

Reviewed By: wickedfoo

Differential Revision: D32646769

fbshipit-source-id: 04a1659173fd51b130ae45d345176b72183cae40
This commit is contained in:
Matthijs Douze
2021-12-07 01:31:07 -08:00
committed by Facebook GitHub Bot
parent 253c7a6070
commit c0052c1533
32 changed files with 265 additions and 397 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c_api/IndexFlat_c.cpp
+3 -3
View File
@@ -47,10 +47,10 @@ int faiss_IndexFlat_new_with(
}
void faiss_IndexFlat_xb(FaissIndexFlat* index, float** p_xb, size_t* p_size) {
auto& xb = reinterpret_cast<IndexFlat*>(index)->xb;
*p_xb = xb.data();
IndexFlat* indexf = reinterpret_cast<IndexFlat*>(index);
*p_xb = indexf->get_xb();
if (p_size) {
*p_size = xb.size();
*p_size = indexf->codes.size() / sizeof(float);
}
}
c_api/IndexLSH_c.cpp
+1 -1
View File
@@ -19,7 +19,7 @@ DEFINE_DESTRUCTOR(IndexLSH)
DEFINE_INDEX_DOWNCAST(IndexLSH)
DEFINE_GETTER(IndexLSH, int, nbits)
DEFINE_GETTER(IndexLSH, int, bytes_per_vec)
DEFINE_GETTER(IndexLSH, int, code_size)
DEFINE_GETTER_PERMISSIVE(IndexLSH, int, rotate_data)
DEFINE_GETTER_PERMISSIVE(IndexLSH, int, train_thresholds)
c_api/IndexLSH_c.h
+1 -1
View File
@@ -25,7 +25,7 @@ FAISS_DECLARE_DESTRUCTOR(IndexLSH)
FAISS_DECLARE_INDEX_DOWNCAST(IndexLSH)
FAISS_DECLARE_GETTER(IndexLSH, int, nbits)
FAISS_DECLARE_GETTER(IndexLSH, int, bytes_per_vec)
FAISS_DECLARE_GETTER(IndexLSH, int, code_size)
FAISS_DECLARE_GETTER(IndexLSH, int, rotate_data)
FAISS_DECLARE_GETTER(IndexLSH, int, train_thresholds)
contrib/inspect_tools.py
+7
View File
@@ -6,6 +6,7 @@
import numpy as np
import faiss
def get_invlist(invlists, l):
""" returns the inverted lists content as a pair of (list_ids, list_codes).
The codes are reshaped to a proper size
@@ -75,3 +76,9 @@ def get_additive_quantizer_codebooks(aq):
codebooks[co[i]:co[i + 1]]
for i in range(aq.M)
]
def get_flat_data(index):
""" copy and return the data matrix in an IndexFlat """
xb = faiss.vector_to_array(index.codes).view("float32")
return xb.reshape(index.ntotal, index.d)
faiss/CMakeLists.txt
+2
View File
@@ -18,6 +18,7 @@ set(FAISS_SRC
IndexBinaryHash.cpp
IndexBinaryIVF.cpp
IndexFlat.cpp
IndexFlatCodes.cpp
IndexHNSW.cpp
IndexIVF.cpp
IndexIVFAdditiveQuantizer.cpp
@@ -93,6 +94,7 @@ set(FAISS_HEADERS
IndexBinaryHash.h
IndexBinaryIVF.h
IndexFlat.h
IndexFlatCodes.h
IndexHNSW.h
IndexIVF.h
IndexIVFAdditiveQuantizer.h
faiss/Index2Layer.cpp
+19 -90
View File
@@ -30,16 +30,6 @@
#include <faiss/utils/distances.h>
#include <faiss/utils/utils.h>
/*
#include <faiss/utils/Heap.h>
#include <faiss/Clustering.h>
#include <faiss/utils/hamming.h>
*/
namespace faiss {
/*************************************
@@ -52,7 +42,7 @@ Index2Layer::Index2Layer(
int M,
int nbit,
MetricType metric)
: Index(quantizer->d, metric),
: IndexFlatCodes(0, quantizer->d, metric),
q1(quantizer, nlist),
pq(quantizer->d, M, nbit) {
is_trained = false;
@@ -116,54 +106,6 @@ void Index2Layer::train(idx_t n, const float* x) {
is_trained = true;
}
void Index2Layer::add(idx_t n, const float* x) {
idx_t bs = 32768;
if (n > bs) {
for (idx_t i0 = 0; i0 < n; i0 += bs) {
idx_t i1 = std::min(i0 + bs, n);
if (verbose) {
printf("Index2Layer::add: adding %" PRId64 ":%" PRId64
" / %" PRId64 "\n",
i0,
i1,
n);
}
add(i1 - i0, x + i0 * d);
}
return;
}
std::vector<idx_t> codes1(n);
q1.quantizer->assign(n, x, codes1.data());
std::vector<float> residuals(n * d);
for (idx_t i = 0; i < n; i++) {
q1.quantizer->compute_residual(
x + i * d, residuals.data() + i * d, codes1[i]);
}
std::vector<uint8_t> codes2(n * code_size_2);
pq.compute_codes(residuals.data(), codes2.data(), n);
codes.resize((ntotal + n) * code_size);
uint8_t* wp = &codes[ntotal * code_size];
{
int i = 0x11223344;
const char* ip = (char*)&i;
FAISS_THROW_IF_NOT_MSG(
ip[0] == 0x44, "works only on a little-endian CPU");
}
// copy to output table
for (idx_t i = 0; i < n; i++) {
memcpy(wp, &codes1[i], code_size_1);
wp += code_size_1;
memcpy(wp, &codes2[i * code_size_2], code_size_2);
wp += code_size_2;
}
ntotal += n;
}
void Index2Layer::search(
idx_t /*n*/,
@@ -174,24 +116,6 @@ void Index2Layer::search(
FAISS_THROW_MSG("not implemented");
}
void Index2Layer::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
std::vector<float> recons1(d);
FAISS_THROW_IF_NOT(i0 >= 0 && i0 + ni <= ntotal);
const uint8_t* rp = &codes[i0 * code_size];
for (idx_t i = 0; i < ni; i++) {
idx_t key = 0;
memcpy(&key, rp, code_size_1);
q1.quantizer->reconstruct(key, recons1.data());
rp += code_size_1;
pq.decode(rp, recons);
for (idx_t j = 0; j < d; j++) {
recons[j] += recons1[j];
}
rp += code_size_2;
recons += d;
}
}
void Index2Layer::transfer_to_IVFPQ(IndexIVFPQ& other) const {
FAISS_THROW_IF_NOT(other.nlist == q1.nlist);
@@ -211,15 +135,6 @@ void Index2Layer::transfer_to_IVFPQ(IndexIVFPQ& other) const {
other.ntotal = ntotal;
}
void Index2Layer::reconstruct(idx_t key, float* recons) const {
reconstruct_n(key, 1, recons);
}
void Index2Layer::reset() {
ntotal = 0;
codes.clear();
}
namespace {
struct Distance2Level : DistanceComputer {
@@ -259,7 +174,7 @@ struct DistanceXPQ4 : Distance2Level {
FAISS_ASSERT(quantizer);
M = storage.pq.M;
pq_l1_tab = quantizer->xb.data();
pq_l1_tab = quantizer->get_xb();
}
float operator()(idx_t i) override {
@@ -368,12 +283,26 @@ DistanceComputer* Index2Layer::get_distance_computer() const {
}
/* The standalone codec interface */
size_t Index2Layer::sa_code_size() const {
return code_size;
}
void Index2Layer::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
FAISS_THROW_IF_NOT(is_trained);
idx_t bs = 32768;
if (n > bs) {
for (idx_t i0 = 0; i0 < n; i0 += bs) {
idx_t i1 = std::min(i0 + bs, n);
if (verbose) {
printf("Index2Layer::add: adding %" PRId64 ":%" PRId64
" / %" PRId64 "\n",
i0,
i1,
n);
}
sa_encode(i1 - i0, x + i0 * d, bytes + i0 * code_size);
}
return;
}
std::unique_ptr<int64_t[]> list_nos(new int64_t[n]);
q1.quantizer->assign(n, x, list_nos.get());
std::vector<float> residuals(n * d);
faiss/Index2Layer.h
+2 -16
View File
@@ -11,6 +11,7 @@
#include <vector>
#include <faiss/IndexFlatCodes.h>
#include <faiss/IndexIVF.h>
#include <faiss/IndexPQ.h>
@@ -24,25 +25,19 @@ struct IndexIVFPQ;
* The class is mainly inteded to store encoded vectors that can be
* accessed randomly, the search function is not implemented.
*/
struct Index2Layer : Index {
struct Index2Layer : IndexFlatCodes {
/// first level quantizer
Level1Quantizer q1;
/// second level quantizer is always a PQ
ProductQuantizer pq;
/// Codes. Size ntotal * code_size.
std::vector<uint8_t> codes;
/// size of the code for the first level (ceil(log8(q1.nlist)))
size_t code_size_1;
/// size of the code for the second level
size_t code_size_2;
/// code_size_1 + code_size_2
size_t code_size;
Index2Layer(
Index* quantizer,
size_t nlist,
@@ -55,8 +50,6 @@ struct Index2Layer : Index {
void train(idx_t n, const float* x) override;
void add(idx_t n, const float* x) override;
/// not implemented
void search(
idx_t n,
@@ -65,19 +58,12 @@ struct Index2Layer : Index {
float* distances,
idx_t* labels) const override;
void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
void reconstruct(idx_t key, float* recons) const override;
void reset() override;
DistanceComputer* get_distance_computer() const override;
/// transfer the flat codes to an IVFPQ index
void transfer_to_IVFPQ(IndexIVFPQ& other) const;
/* The standalone codec interface */
size_t sa_code_size() const override;
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
};
faiss/IndexAdditiveQuantizer.cpp
+1 -17
View File
@@ -32,19 +32,12 @@ IndexAdditiveQuantizer::IndexAdditiveQuantizer(
idx_t d,
AdditiveQuantizer* aq,
MetricType metric):
Index(d, metric), aq(aq)
IndexFlatCodes(aq->code_size, d, metric), aq(aq)
{
FAISS_THROW_IF_NOT(metric == METRIC_INNER_PRODUCT || metric == METRIC_L2);
}
void IndexAdditiveQuantizer::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT(is_trained);
codes.resize((n + ntotal) * aq->code_size);
aq->compute_codes(x, &codes[ntotal * aq->code_size], n);
ntotal += n;
}
namespace {
template <class VectorDistance, class ResultHandler>
@@ -163,15 +156,6 @@ void IndexAdditiveQuantizer::search(
}
}
void IndexAdditiveQuantizer::reset() {
codes.clear();
ntotal = 0;
}
size_t IndexAdditiveQuantizer::sa_code_size() const {
return code_size;
}
void IndexAdditiveQuantizer::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
return aq->compute_codes(x, bytes, n);
}
faiss/IndexAdditiveQuantizer.h
+2 -14
View File
@@ -13,7 +13,7 @@
#include <cstdint>
#include <vector>
#include <faiss/Index.h>
#include <faiss/IndexFlatCodes.h>
#include <faiss/impl/LocalSearchQuantizer.h>
#include <faiss/impl/ResidualQuantizer.h>
#include <faiss/impl/platform_macros.h>
@@ -21,7 +21,7 @@
namespace faiss {
/// Abstract class for additive quantizers. The search functions are in common.
struct IndexAdditiveQuantizer : Index {
struct IndexAdditiveQuantizer : IndexFlatCodes {
// the quantizer, this points to the relevant field in the inheriting
// classes
AdditiveQuantizer* aq;
@@ -32,12 +32,6 @@ struct IndexAdditiveQuantizer : Index {
AdditiveQuantizer* aq = nullptr,
MetricType metric = METRIC_L2);
/// size of residual quantizer codes + norms
size_t code_size;
/// Codes. Size ntotal * rq.code_size
std::vector<uint8_t> codes;
void search(
idx_t n,
const float* x,
@@ -45,13 +39,7 @@ struct IndexAdditiveQuantizer : Index {
float* distances,
idx_t* labels) const override;
void reset() override;
void add(idx_t n, const float* x) override;
/* The standalone codec interface */
size_t sa_code_size() const override;
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
faiss/IndexFlat.cpp
+22 -52
View File
@@ -19,17 +19,8 @@
namespace faiss {
IndexFlat::IndexFlat(idx_t d, MetricType metric) : Index(d, metric) {}
void IndexFlat::add(idx_t n, const float* x) {
xb.insert(xb.end(), x, x + n * d);
ntotal += n;
}
void IndexFlat::reset() {
xb.clear();
ntotal = 0;
}
IndexFlat::IndexFlat(idx_t d, MetricType metric)
: IndexFlatCodes(sizeof(float) * d, d, metric) {}
void IndexFlat::search(
idx_t n,
@@ -43,14 +34,14 @@ void IndexFlat::search(
if (metric_type == METRIC_INNER_PRODUCT) {
float_minheap_array_t res = {size_t(n), size_t(k), labels, distances};
knn_inner_product(x, xb.data(), d, n, ntotal, &res);
knn_inner_product(x, get_xb(), d, n, ntotal, &res);
} else if (metric_type == METRIC_L2) {
float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
knn_L2sqr(x, xb.data(), d, n, ntotal, &res);
knn_L2sqr(x, get_xb(), d, n, ntotal, &res);
} else {
float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances};
knn_extra_metrics(
x, xb.data(), d, n, ntotal, metric_type, metric_arg, &res);
x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res);
}
}
@@ -62,10 +53,10 @@ void IndexFlat::range_search(
switch (metric_type) {
case METRIC_INNER_PRODUCT:
range_search_inner_product(
x, xb.data(), d, n, ntotal, radius, result);
x, get_xb(), d, n, ntotal, radius, result);
break;
case METRIC_L2:
range_search_L2sqr(x, xb.data(), d, n, ntotal, radius, result);
range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result);
break;
default:
FAISS_THROW_MSG("metric type not supported");
@@ -80,37 +71,16 @@ void IndexFlat::compute_distance_subset(
const idx_t* labels) const {
switch (metric_type) {
case METRIC_INNER_PRODUCT:
fvec_inner_products_by_idx(
distances, x, xb.data(), labels, d, n, k);
fvec_inner_products_by_idx(distances, x, get_xb(), labels, d, n, k);
break;
case METRIC_L2:
fvec_L2sqr_by_idx(distances, x, xb.data(), labels, d, n, k);
fvec_L2sqr_by_idx(distances, x, get_xb(), labels, d, n, k);
break;
default:
FAISS_THROW_MSG("metric type not supported");
}
}
size_t IndexFlat::remove_ids(const IDSelector& sel) {
idx_t j = 0;
for (idx_t i = 0; i < ntotal; i++) {
if (sel.is_member(i)) {
// should be removed
} else {
if (i > j) {
memmove(&xb[d * j], &xb[d * i], sizeof(xb[0]) * d);
}
j++;
}
}
size_t nremove = ntotal - j;
if (nremove > 0) {
ntotal = j;
xb.resize(ntotal * d);
}
return nremove;
}
namespace {
struct FlatL2Dis : DistanceComputer {
@@ -133,7 +103,7 @@ struct FlatL2Dis : DistanceComputer {
: d(storage.d),
nb(storage.ntotal),
q(q),
b(storage.xb.data()),
b(storage.get_xb()),
ndis(0) {}
void set_query(const float* x) override {
@@ -161,7 +131,7 @@ struct FlatIPDis : DistanceComputer {
: d(storage.d),
nb(storage.ntotal),
q(q),
b(storage.xb.data()),
b(storage.get_xb()),
ndis(0) {}
void set_query(const float* x) override {
@@ -178,25 +148,24 @@ DistanceComputer* IndexFlat::get_distance_computer() const {
return new FlatIPDis(*this);
} else {
return get_extra_distance_computer(
d, metric_type, metric_arg, ntotal, xb.data());
d, metric_type, metric_arg, ntotal, get_xb());
}
}
void IndexFlat::reconstruct(idx_t key, float* recons) const {
memcpy(recons, &(xb[key * d]), sizeof(*recons) * d);
}
/* The standalone codec interface */
size_t IndexFlat::sa_code_size() const {
return sizeof(float) * d;
memcpy(recons, &(codes[key * code_size]), code_size);
}
void IndexFlat::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
memcpy(bytes, x, sizeof(float) * d * n);
if (n > 0) {
memcpy(bytes, x, sizeof(float) * d * n);
}
}
void IndexFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
memcpy(x, bytes, sizeof(float) * d * n);
if (n > 0) {
memcpy(x, bytes, sizeof(float) * d * n);
}
}
/***************************************************
@@ -211,9 +180,9 @@ IndexFlat1D::IndexFlat1D(bool continuous_update)
void IndexFlat1D::update_permutation() {
perm.resize(ntotal);
if (ntotal < 1000000) {
fvec_argsort(ntotal, xb.data(), (size_t*)perm.data());
fvec_argsort(ntotal, get_xb(), (size_t*)perm.data());
} else {
fvec_argsort_parallel(ntotal, xb.data(), (size_t*)perm.data());
fvec_argsort_parallel(ntotal, get_xb(), (size_t*)perm.data());
}
}
@@ -238,6 +207,7 @@ void IndexFlat1D::search(
FAISS_THROW_IF_NOT_MSG(
perm.size() == ntotal, "Call update_permutation before search");
const float* xb = get_xb();
#pragma omp parallel for
for (idx_t i = 0; i < n; i++) {
faiss/IndexFlat.h
+9 -15
View File
@@ -12,21 +12,14 @@
#include <vector>
#include <faiss/Index.h>
#include <faiss/IndexFlatCodes.h>
namespace faiss {
/** Index that stores the full vectors and performs exhaustive search */
struct IndexFlat : Index {
/// database vectors, size ntotal * d
std::vector<float> xb;
struct IndexFlat : IndexFlatCodes {
explicit IndexFlat(idx_t d, MetricType metric = METRIC_L2);
void add(idx_t n, const float* x) override;
void reset() override;
void search(
idx_t n,
const float* x,
@@ -57,18 +50,19 @@ struct IndexFlat : Index {
float* distances,
const idx_t* labels) const;
/** remove some ids. NB that Because of the structure of the
* indexing structure, the semantics of this operation are
* different from the usual ones: the new ids are shifted */
size_t remove_ids(const IDSelector& sel) override;
// get pointer to the floating point data
float* get_xb() {
return (float*)codes.data();
}
const float* get_xb() const {
return (const float*)codes.data();
}
IndexFlat() {}
DistanceComputer* get_distance_computer() const override;
/* The stanadlone codec interface (just memcopies in this case) */
size_t sa_code_size() const override;
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
faiss/IndexFlatCodes.cpp
+67
View File
@@ -0,0 +1,67 @@
/**
* Copyright (c) Facebook, Inc. and its affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <faiss/IndexFlatCodes.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
namespace faiss {
IndexFlatCodes::IndexFlatCodes(size_t code_size, idx_t d, MetricType metric)
: Index(d, metric), code_size(code_size) {}
IndexFlatCodes::IndexFlatCodes() : code_size(0) {}
void IndexFlatCodes::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT(is_trained);
codes.resize((ntotal + n) * code_size);
sa_encode(n, x, &codes[ntotal * code_size]);
ntotal += n;
}
void IndexFlatCodes::reset() {
codes.clear();
ntotal = 0;
}
size_t IndexFlatCodes::sa_code_size() const {
return code_size;
}
size_t IndexFlatCodes::remove_ids(const IDSelector& sel) {
idx_t j = 0;
for (idx_t i = 0; i < ntotal; i++) {
if (sel.is_member(i)) {
// should be removed
} else {
if (i > j) {
memmove(&codes[code_size * j],
&codes[code_size * i],
code_size);
}
j++;
}
}
size_t nremove = ntotal - j;
if (nremove > 0) {
ntotal = j;
codes.resize(ntotal * code_size);
}
return nremove;
}
void IndexFlatCodes::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
sa_decode(ni, codes.data() + i0 * code_size, recons);
}
void IndexFlatCodes::reconstruct(idx_t key, float* recons) const {
reconstruct_n(key, 1, recons);
}
} // namespace faiss
faiss/IndexFlatCodes.h
+47
View File
@@ -0,0 +1,47 @@
/**
* Copyright (c) Facebook, Inc. and its affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
// -*- c++ -*-
#pragma once
#include <faiss/Index.h>
#include <vector>
namespace faiss {
/** Index that encodes all vectors as fixed-size codes (size code_size). Storage
* is in the codes vector */
struct IndexFlatCodes : Index {
size_t code_size;
/// encoded dataset, size ntotal * code_size
std::vector<uint8_t> codes;
IndexFlatCodes();
IndexFlatCodes(size_t code_size, idx_t d, MetricType metric = METRIC_L2);
/// default add uses sa_encode
void add(idx_t n, const float* x) override;
void reset() override;
/// reconstruction using the codec interface
void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
void reconstruct(idx_t key, float* recons) const override;
size_t sa_code_size() const override;
/** remove some ids. NB that Because of the structure of the
* indexing structure, the semantics of this operation are
* different from the usual ones: the new ids are shifted */
size_t remove_ids(const IDSelector& sel) override;
};
} // namespace faiss
faiss/IndexLSH.cpp
+4 -30
View File
@@ -5,8 +5,6 @@
* LICENSE file in the root directory of this source tree.
*/
// -*- c++ -*-
#include <faiss/IndexLSH.h>
#include <cstdio>
@@ -25,15 +23,13 @@ namespace faiss {
***************************************************************/
IndexLSH::IndexLSH(idx_t d, int nbits, bool rotate_data, bool train_thresholds)
: Index(d),
: IndexFlatCodes((nbits + 7) / 8, d),
nbits(nbits),
rotate_data(rotate_data),
train_thresholds(train_thresholds),
rrot(d, nbits) {
is_trained = !train_thresholds;
bytes_per_vec = (nbits + 7) / 8;
if (rotate_data) {
rrot.init(5);
} else {
@@ -41,11 +37,7 @@ IndexLSH::IndexLSH(idx_t d, int nbits, bool rotate_data, bool train_thresholds)
}
}
IndexLSH::IndexLSH()
: nbits(0),
bytes_per_vec(0),
rotate_data(false),
train_thresholds(false) {}
IndexLSH::IndexLSH() : nbits(0), rotate_data(false), train_thresholds(false) {}
const float* IndexLSH::apply_preprocess(idx_t n, const float* x) const {
float* xt = nullptr;
@@ -106,15 +98,6 @@ void IndexLSH::train(idx_t n, const float* x) {
is_trained = true;
}
void IndexLSH::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT(is_trained);
codes.resize((ntotal + n) * bytes_per_vec);
sa_encode(n, x, &codes[ntotal * bytes_per_vec]);
ntotal += n;
}
void IndexLSH::search(
idx_t n,
const float* x,
@@ -127,7 +110,7 @@ void IndexLSH::search(
const float* xt = apply_preprocess(n, x);
ScopeDeleter<float> del(xt == x ? nullptr : xt);
uint8_t* qcodes = new uint8_t[n * bytes_per_vec];
uint8_t* qcodes = new uint8_t[n * code_size];
ScopeDeleter<uint8_t> del2(qcodes);
fvecs2bitvecs(xt, qcodes, nbits, n);
@@ -137,7 +120,7 @@ void IndexLSH::search(
int_maxheap_array_t res = {size_t(n), size_t(k), labels, idistances};
hammings_knn_hc(&res, qcodes, codes.data(), ntotal, bytes_per_vec, true);
hammings_knn_hc(&res, qcodes, codes.data(), ntotal, code_size, true);
// convert distances to floats
for (int i = 0; i < k * n; i++)
@@ -158,15 +141,6 @@ void IndexLSH::transfer_thresholds(LinearTransform* vt) {
thresholds.clear();
}
void IndexLSH::reset() {
codes.clear();
ntotal = 0;
}
size_t IndexLSH::sa_code_size() const {
return bytes_per_vec;
}
void IndexLSH::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
FAISS_THROW_IF_NOT(is_trained);
const float* xt = apply_preprocess(n, x);
faiss/IndexLSH.h
+2 -14
View File
@@ -12,17 +12,14 @@
#include <vector>
#include <faiss/Index.h>
#include <faiss/IndexFlatCodes.h>
#include <faiss/VectorTransform.h>
namespace faiss {
/** The sign of each vector component is put in a binary signature */
struct IndexLSH : Index {
typedef unsigned char uint8_t;
struct IndexLSH : IndexFlatCodes {
int nbits; ///< nb of bits per vector
int bytes_per_vec; ///< nb of 8-bits per encoded vector
bool rotate_data; ///< whether to apply a random rotation to input
bool train_thresholds; ///< whether we train thresholds or use 0
@@ -30,9 +27,6 @@ struct IndexLSH : Index {
std::vector<float> thresholds; ///< thresholds to compare with
/// encoded dataset
std::vector<uint8_t> codes;
IndexLSH(
idx_t d,
int nbits,
@@ -50,8 +44,6 @@ struct IndexLSH : Index {
void train(idx_t n, const float* x) override;
void add(idx_t n, const float* x) override;
void search(
idx_t n,
const float* x,
@@ -59,7 +51,6 @@ struct IndexLSH : Index {
float* distances,
idx_t* labels) const override;
void reset() override;
/// transfer the thresholds to a pre-processing stage (and unset
/// train_thresholds)
@@ -72,9 +63,6 @@ struct IndexLSH : Index {
/* standalone codec interface.
*
* The vectors are decoded to +/- 1 (not 0, 1) */
size_t sa_code_size() const override;
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
faiss/IndexPQ.cpp
+2 -51
View File
@@ -28,12 +28,13 @@ namespace faiss {
********************************************************/
IndexPQ::IndexPQ(int d, size_t M, size_t nbits, MetricType metric)
: Index(d, metric), pq(d, M, nbits) {
: IndexFlatCodes(0, d, metric), pq(d, M, nbits) {
is_trained = false;
do_polysemous_training = false;
polysemous_ht = nbits * M + 1;
search_type = ST_PQ;
encode_signs = false;
code_size = pq.code_size;
}
IndexPQ::IndexPQ() {
@@ -69,53 +70,6 @@ void IndexPQ::train(idx_t n, const float* x) {
is_trained = true;
}
void IndexPQ::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT(is_trained);
codes.resize((n + ntotal) * pq.code_size);
pq.compute_codes(x, &codes[ntotal * pq.code_size], n);
ntotal += n;
}
size_t IndexPQ::remove_ids(const IDSelector& sel) {
idx_t j = 0;
for (idx_t i = 0; i < ntotal; i++) {
if (sel.is_member(i)) {
// should be removed
} else {
if (i > j) {
memmove(&codes[pq.code_size * j],
&codes[pq.code_size * i],
pq.code_size);
}
j++;
}
}
size_t nremove = ntotal - j;
if (nremove > 0) {
ntotal = j;
codes.resize(ntotal * pq.code_size);
}
return nremove;
}
void IndexPQ::reset() {
codes.clear();
ntotal = 0;
}
void IndexPQ::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
FAISS_THROW_IF_NOT(ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
for (idx_t i = 0; i < ni; i++) {
const uint8_t* code = &codes[(i0 + i) * pq.code_size];
pq.decode(code, recons + i * d);
}
}
void IndexPQ::reconstruct(idx_t key, float* recons) const {
FAISS_THROW_IF_NOT(key >= 0 && key < ntotal);
pq.decode(&codes[key * pq.code_size], recons);
}
namespace {
template <class PQDecoder>
@@ -457,9 +411,6 @@ void IndexPQ::search_core_polysemous(
}
/* The standalone codec interface (just remaps to the PQ functions) */
size_t IndexPQ::sa_code_size() const {
return pq.code_size;
}
void IndexPQ::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
pq.compute_codes(x, bytes, n);
faiss/IndexPQ.h
+2 -17
View File
@@ -12,7 +12,7 @@
#include <vector>
#include <faiss/Index.h>
#include <faiss/IndexFlatCodes.h>
#include <faiss/impl/PolysemousTraining.h>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/impl/platform_macros.h>
@@ -21,13 +21,10 @@ namespace faiss {
/** Index based on a product quantizer. Stored vectors are
* approximated by PQ codes. */
struct IndexPQ : Index {
struct IndexPQ : IndexFlatCodes {
/// The product quantizer used to encode the vectors
ProductQuantizer pq;
/// Codes. Size ntotal * pq.code_size
std::vector<uint8_t> codes;
/** Constructor.
*
* @param d dimensionality of the input vectors
@@ -43,8 +40,6 @@ struct IndexPQ : Index {
void train(idx_t n, const float* x) override;
void add(idx_t n, const float* x) override;
void search(
idx_t n,
const float* x,
@@ -52,17 +47,7 @@ struct IndexPQ : Index {
float* distances,
idx_t* labels) const override;
void reset() override;
void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
void reconstruct(idx_t key, float* recons) const override;
size_t remove_ids(const IDSelector& sel) override;
/* The standalone codec interface */
size_t sa_code_size() const override;
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
faiss/IndexScalarQuantizer.cpp
+1 -28
View File
@@ -29,7 +29,7 @@ IndexScalarQuantizer::IndexScalarQuantizer(
int d,
ScalarQuantizer::QuantizerType qtype,
MetricType metric)
: Index(d, metric), sq(d, qtype) {
: IndexFlatCodes(0, d, metric), sq(d, qtype) {
is_trained = qtype == ScalarQuantizer::QT_fp16 ||
qtype == ScalarQuantizer::QT_8bit_direct;
code_size = sq.code_size;
@@ -43,13 +43,6 @@ void IndexScalarQuantizer::train(idx_t n, const float* x) {
is_trained = true;
}
void IndexScalarQuantizer::add(idx_t n, const float* x) {
FAISS_THROW_IF_NOT(is_trained);
codes.resize((n + ntotal) * code_size);
sq.compute_codes(x, &codes[ntotal * code_size], n);
ntotal += n;
}
void IndexScalarQuantizer::search(
idx_t n,
const float* x,
@@ -100,27 +93,7 @@ DistanceComputer* IndexScalarQuantizer::get_distance_computer() const {
return dc;
}
void IndexScalarQuantizer::reset() {
codes.clear();
ntotal = 0;
}
void IndexScalarQuantizer::reconstruct_n(idx_t i0, idx_t ni, float* recons)
const {
std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
for (size_t i = 0; i < ni; i++) {
squant->decode_vector(&codes[(i + i0) * code_size], recons + i * d);
}
}
void IndexScalarQuantizer::reconstruct(idx_t key, float* recons) const {
reconstruct_n(key, 1, recons);
}
/* Codec interface */
size_t IndexScalarQuantizer::sa_code_size() const {
return sq.code_size;
}
void IndexScalarQuantizer::sa_encode(idx_t n, const float* x, uint8_t* bytes)
const {
faiss/IndexScalarQuantizer.h
+2 -16
View File
@@ -13,6 +13,7 @@
#include <stdint.h>
#include <vector>
#include <faiss/IndexFlatCodes.h>
#include <faiss/IndexIVF.h>
#include <faiss/impl/ScalarQuantizer.h>
@@ -24,15 +25,10 @@ namespace faiss {
* (default).
*/
struct IndexScalarQuantizer : Index {
struct IndexScalarQuantizer : IndexFlatCodes {
/// Used to encode the vectors
ScalarQuantizer sq;
/// Codes. Size ntotal * pq.code_size
std::vector<uint8_t> codes;
size_t code_size;
/** Constructor.
*
* @param d dimensionality of the input vectors
@@ -48,8 +44,6 @@ struct IndexScalarQuantizer : Index {
void train(idx_t n, const float* x) override;
void add(idx_t n, const float* x) override;
void search(
idx_t n,
const float* x,
@@ -57,17 +51,9 @@ struct IndexScalarQuantizer : Index {
float* distances,
idx_t* labels) const override;
void reset() override;
void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
void reconstruct(idx_t key, float* recons) const override;
DistanceComputer* get_distance_computer() const override;
/* standalone codec interface */
size_t sa_code_size() const override;
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
faiss/gpu/GpuCloner.cpp
+2 -2
View File
@@ -40,7 +40,7 @@ void ToCPUCloner::merge_index(Index* dst, Index* src, bool successive_ids) {
auto ifl2 = dynamic_cast<const IndexFlat*>(src);
FAISS_ASSERT(ifl2);
FAISS_ASSERT(successive_ids);
ifl->add(ifl2->ntotal, ifl2->xb.data());
ifl->add(ifl2->ntotal, ifl2->get_xb());
} else if (auto ifl = dynamic_cast<IndexIVFFlat*>(dst)) {
auto ifl2 = dynamic_cast<IndexIVFFlat*>(src);
FAISS_ASSERT(ifl2);
@@ -329,7 +329,7 @@ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
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);
shards[i]->add(i1 - i0, index_flat->get_xb() + i0 * index->d);
}
}
}
faiss/gpu/GpuIndexFlat.cu
+5 -4
View File
@@ -118,7 +118,7 @@ void GpuIndexFlat::copyFrom(const faiss::IndexFlat* index) {
// The index could be empty
if (index->ntotal > 0) {
data_->add(
index->xb.data(),
index->get_xb(),
index->ntotal,
resources_->getDefaultStream(config_.device));
}
@@ -128,19 +128,20 @@ void GpuIndexFlat::copyTo(faiss::IndexFlat* index) const {
DeviceScope scope(config_.device);
GpuIndex::copyTo(index);
index->code_size = sizeof(float) * this->d;
FAISS_ASSERT(data_);
FAISS_ASSERT(data_->getSize() == this->ntotal);
index->xb.resize(this->ntotal * this->d);
index->codes.resize(this->ntotal * index->code_size);
auto stream = resources_->getDefaultStream(config_.device);
if (this->ntotal > 0) {
if (flatConfig_.useFloat16) {
auto vecFloat32 = data_->getVectorsFloat32Copy(stream);
fromDevice(vecFloat32, index->xb.data(), stream);
fromDevice(vecFloat32, index->get_xb(), stream);
} else {
fromDevice(data_->getVectorsFloat32Ref(), index->xb.data(), stream);
fromDevice(data_->getVectorsFloat32Ref(), index->get_xb(), stream);
}
}
}
faiss/impl/AdditiveQuantizer.cpp
+1 -1
View File
@@ -137,7 +137,7 @@ uint32_t AdditiveQuantizer::encode_qcint(float x) const {
}
float AdditiveQuantizer::decode_qcint(uint32_t c) const {
return qnorm.xb[c];
return qnorm.get_xb()[c];
}
void AdditiveQuantizer::pack_codes(
faiss/impl/index_read.cpp
+14 -8
View File
@@ -146,9 +146,10 @@ VectorTransform* read_VectorTransform(IOReader* f) {
vt = itqt;
} else {
FAISS_THROW_FMT(
"fourcc %ud (\"%s\") not recognized",
"fourcc %ud (\"%s\") not recognized in %s",
h,
fourcc_inv_printable(h).c_str());
fourcc_inv_printable(h).c_str(),
f->name.c_str());
}
READ1(vt->d_in);
READ1(vt->d_out);
@@ -271,7 +272,7 @@ static void read_AdditiveQuantizer(AdditiveQuantizer* aq, IOReader* f) {
READ1(aq->norm_max);
if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 ||
aq->search_type == AdditiveQuantizer::ST_norm_cqint4) {
READVECTOR(aq->qnorm.xb);
READXBVECTOR(aq->qnorm.codes);
}
aq->set_derived_values();
}
@@ -472,8 +473,10 @@ Index* read_index(IOReader* f, int io_flags) {
idxf = new IndexFlat();
}
read_index_header(idxf, f);
READVECTOR(idxf->xb);
FAISS_THROW_IF_NOT(idxf->xb.size() == idxf->ntotal * idxf->d);
idxf->code_size = idxf->d * sizeof(float);
READXBVECTOR(idxf->codes);
FAISS_THROW_IF_NOT(
idxf->codes.size() == idxf->ntotal * idxf->code_size);
// leak!
idx = idxf;
} else if (h == fourcc("IxHE") || h == fourcc("IxHe")) {
@@ -483,7 +486,9 @@ Index* read_index(IOReader* f, int io_flags) {
READ1(idxl->rotate_data);
READ1(idxl->train_thresholds);
READVECTOR(idxl->thresholds);
READ1(idxl->bytes_per_vec);
int code_size_i;
READ1(code_size_i);
idxl->code_size = code_size_i;
if (h == fourcc("IxHE")) {
FAISS_THROW_IF_NOT_FMT(
idxl->nbits % 64 == 0,
@@ -491,7 +496,7 @@ Index* read_index(IOReader* f, int io_flags) {
"nbits multiple of 64 (got %d)",
(int)idxl->nbits);
// leak
idxl->bytes_per_vec *= 8;
idxl->code_size *= 8;
}
{
RandomRotationMatrix* rrot = dynamic_cast<RandomRotationMatrix*>(
@@ -504,7 +509,7 @@ Index* read_index(IOReader* f, int io_flags) {
FAISS_THROW_IF_NOT(
idxl->rrot.d_in == idxl->d && idxl->rrot.d_out == idxl->nbits);
FAISS_THROW_IF_NOT(
idxl->codes.size() == idxl->ntotal * idxl->bytes_per_vec);
idxl->codes.size() == idxl->ntotal * idxl->code_size);
idx = idxl;
} else if (
h == fourcc("IxPQ") || h == fourcc("IxPo") || h == fourcc("IxPq")) {
@@ -512,6 +517,7 @@ Index* read_index(IOReader* f, int io_flags) {
IndexPQ* idxp = new IndexPQ();
read_index_header(idxp, f);
read_ProductQuantizer(&idxp->pq, f);
idxp->code_size = idxp->pq.code_size;
READVECTOR(idxp->codes);
if (h == fourcc("IxPo") || h == fourcc("IxPq")) {
READ1(idxp->search_type);
faiss/impl/index_write.cpp
+4 -3
View File
@@ -171,7 +171,7 @@ static void write_AdditiveQuantizer(const AdditiveQuantizer* aq, IOWriter* f) {
WRITE1(aq->norm_max);
if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 ||
aq->search_type == AdditiveQuantizer::ST_norm_cqint4) {
WRITEVECTOR(aq->qnorm.xb);
WRITEXBVECTOR(aq->qnorm.codes);
}
}
@@ -345,7 +345,7 @@ void write_index(const Index* idx, IOWriter* f) {
: "IxFl");
WRITE1(h);
write_index_header(idx, f);
WRITEVECTOR(idxf->xb);
WRITEXBVECTOR(idxf->codes);
} else if (const IndexLSH* idxl = dynamic_cast<const IndexLSH*>(idx)) {
uint32_t h = fourcc("IxHe");
WRITE1(h);
@@ -354,7 +354,8 @@ void write_index(const Index* idx, IOWriter* f) {
WRITE1(idxl->rotate_data);
WRITE1(idxl->train_thresholds);
WRITEVECTOR(idxl->thresholds);
WRITE1(idxl->bytes_per_vec);
int code_size_i = idxl->code_size;
WRITE1(code_size_i);
write_VectorTransform(&idxl->rrot, f);
WRITEVECTOR(idxl->codes);
} else if (const IndexPQ* idxp = dynamic_cast<const IndexPQ*>(idx)) {
faiss/impl/io.cpp
+1 -1
View File
@@ -240,7 +240,7 @@ uint32_t fourcc(const std::string& sx) {
void fourcc_inv(uint32_t x, char str[5]) {
*(uint32_t*)str = x;
str[5] = 0;
str[4] = 0;
}
std::string fourcc_inv(uint32_t x) {
faiss/impl/io_macros.h
+20
View File
@@ -66,3 +66,23 @@
WRITEANDCHECK(&size, 1); \
WRITEANDCHECK((vec).data(), size); \
}
// read/write xb vector for backwards compatibility of IndexFlat
#define WRITEXBVECTOR(vec) \
{ \
FAISS_THROW_IF_NOT((vec).size() % 4 == 0); \
size_t size = (vec).size() / 4; \
WRITEANDCHECK(&size, 1); \
WRITEANDCHECK((vec).data(), size * 4); \
}
#define READXBVECTOR(vec) \
{ \
size_t size; \
READANDCHECK(&size, 1); \
FAISS_THROW_IF_NOT(size >= 0 && size < (uint64_t{1} << 40)); \
size *= 4; \
(vec).resize(size); \
READANDCHECK((vec).data(), size); \
}
faiss/python/swigfaiss.swig
+1
View File
@@ -374,6 +374,7 @@ void gpu_sync_all_devices()
%newobject *::get_distance_computer() const;
%include <faiss/Index.h>
%include <faiss/IndexFlatCodes.h>
%include <faiss/IndexFlat.h>
%include <faiss/Clustering.h>
tests/test_build_blocks.py
+1 -1
View File
@@ -75,7 +75,7 @@ class TestRevSwigPtr(unittest.TestCase):
i * 10 + np.array([1, 2, 3, 4], dtype='float32')
for i in range(5)])
index.add(xb0)
xb = faiss.rev_swig_ptr(index.xb.data(), 4 * 5).reshape(5, 4)
xb = faiss.rev_swig_ptr(index.get_xb(), 4 * 5).reshape(5, 4)
self.assertEqual(np.abs(xb0 - xb).sum(), 0)
tests/test_contrib.py
+8 -1
View File
@@ -174,7 +174,6 @@ class TestExhaustiveSearch(unittest.TestCase):
)
class TestInspect(unittest.TestCase):
def test_LinearTransform(self):
@@ -194,6 +193,14 @@ class TestInspect(unittest.TestCase):
ynew = x @ A.T + b
np.testing.assert_array_almost_equal(yref, ynew)
def test_IndexFlat(self):
xb = np.random.rand(13, 20).astype('float32')
index = faiss.IndexFlatL2(20)
index.add(xb)
np.testing.assert_array_equal(
xb, inspect_tools.get_flat_data(index)
)
class TestRangeEval(unittest.TestCase):
tests/test_index_accuracy.py
+2 -1
View File
@@ -192,7 +192,8 @@ class TestSQFlavors(unittest.TestCase):
nlist = index.nlist
quantizer = faiss.downcast_index(index.quantizer)
quantizer2 = faiss.IndexFlat(d2, index.metric_type)
centroids = faiss.vector_to_array(quantizer.xb).reshape(nlist, d)
centroids = faiss.vector_to_array(quantizer.codes)
centroids = centroids.view("float32").reshape(nlist, d)
centroids2 = self.add2columns(centroids)
quantizer2.add(centroids2)
index2 = faiss.IndexIVFScalarQuantizer(
tests/test_index_composite.py
+2 -1
View File
@@ -77,7 +77,8 @@ class TestRemove(unittest.TestCase):
xb[:, 0] = np.arange(10, dtype='int64') + 1000
index.add(xb)
index.remove_ids(np.arange(5, dtype='int64') * 2)
xb2 = faiss.vector_float_to_array(index.xb).reshape(5, 5)
xb2 = faiss.vector_float_to_array(index.codes)
xb2 = xb2.view("float32").reshape(5, 5)
assert np.all(xb2[:, 0] == xb[np.arange(5) * 2 + 1, 0])
def test_remove_id_map(self):
tests/test_io.py
+8 -9
View File
@@ -13,8 +13,6 @@ import sys
import pickle
from multiprocessing.dummy import Pool as ThreadPool
from common_faiss_tests import get_dataset, get_dataset_2
class TestIOVariants(unittest.TestCase):
@@ -77,9 +75,10 @@ class TestCallbacks(unittest.TestCase):
index2 = faiss.deserialize_index(np.frombuffer(buf, dtype='uint8'))
self.assertEqual(index.d, index2.d)
self.assertTrue(np.all(
faiss.vector_to_array(index.xb) == faiss.vector_to_array(index2.xb)
))
np.testing.assert_array_equal(
faiss.vector_to_array(index.codes),
faiss.vector_to_array(index2.codes)
)
# This is not a callable function: shoudl raise an exception
writer = faiss.PyCallbackIOWriter("blabla")
@@ -132,8 +131,8 @@ class TestCallbacks(unittest.TestCase):
self.assertEqual(index.d, index2.d)
np.testing.assert_array_equal(
faiss.vector_to_array(index.xb),
faiss.vector_to_array(index2.xb)
faiss.vector_to_array(index.codes),
faiss.vector_to_array(index2.codes)
)
# This is not a callable function: should raise an exception
@@ -178,8 +177,8 @@ class TestCallbacks(unittest.TestCase):
self.assertEqual(index.d, index2.d)
np.testing.assert_array_equal(
faiss.vector_to_array(index.xb),
faiss.vector_to_array(index2.xb)
faiss.vector_to_array(index.codes),
faiss.vector_to_array(index2.codes)
)
finally: