Skip to content

Commit f7aedbd

Browse files
mdouzemdouze
committed
sync with FB version 2017-07-18
- implemented ScalarQuantizer (without IVF) - implemented update for IndexIVFFlat - implemented L2 normalization preproc
1 parent 602deba commit f7aedbd

24 files changed

Lines changed: 4531 additions & 1888 deletions

‎AutoTune.cpp‎

Lines changed: 56 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
#include "IndexIVF.h"
2323
#include "IndexIVFPQ.h"
2424
#include "MetaIndexes.h"
25-
#include "IndexIVFScalarQuantizer.h"
25+
#include "IndexScalarQuantizer.h"
2626

2727

2828
namespace faiss {
@@ -623,18 +623,28 @@ void ParameterSpace::explore (Index *index,
623623
* index_factory
624624
***************************************************************/
625625

626+
namespace {
626627

628+
struct VTChain {
629+
std::vector<VectorTransform *> chain;
630+
~VTChain () {
631+
for (int i = 0; i < chain.size(); i++) {
632+
delete chain[i];
633+
}
634+
}
635+
};
636+
637+
}
627638

628639
Index *index_factory (int d, const char *description_in, MetricType metric)
629640
{
630-
VectorTransform *vt = nullptr;
641+
VTChain vts;
631642
Index *coarse_quantizer = nullptr;
632643
Index *index = nullptr;
633644
bool add_idmap = false;
634645
bool make_IndexRefineFlat = false;
635646

636647
ScopeDeleter1<Index> del_coarse_quantizer, del_index;
637-
ScopeDeleter1<VectorTransform> del_vt;
638648

639649
char description[strlen(description_in) + 1];
640650
char *ptr;
@@ -656,18 +666,27 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
656666
Index *index_1 = nullptr;
657667

658668
// VectorTransforms
659-
if (!vt && sscanf (tok, "PCA%d", &d_out) == 1) {
669+
if (sscanf (tok, "PCA%d", &d_out) == 1) {
660670
vt_1 = new PCAMatrix (d, d_out);
661671
d = d_out;
662-
} else if (!vt && sscanf (tok, "PCAR%d", &d_out) == 1) {
672+
} else if (sscanf (tok, "PCAR%d", &d_out) == 1) {
663673
vt_1 = new PCAMatrix (d, d_out, 0, true);
664674
d = d_out;
665-
} else if (!vt && sscanf (tok, "OPQ%d_%d", &opq_M, &d_out) == 2) {
675+
} else if (sscanf (tok, "PCAW%d", &d_out) == 1) {
676+
vt_1 = new PCAMatrix (d, d_out, -0.5, false);
677+
d = d_out;
678+
} else if (sscanf (tok, "PCAWR%d", &d_out) == 1) {
679+
vt_1 = new PCAMatrix (d, d_out, -0.5, true);
680+
d = d_out;
681+
} else if (sscanf (tok, "OPQ%d_%d", &opq_M, &d_out) == 2) {
666682
vt_1 = new OPQMatrix (d, opq_M, d_out);
667683
d = d_out;
668-
} else if (!vt && sscanf (tok, "OPQ%d", &opq_M) == 1) {
684+
} else if (sscanf (tok, "OPQ%d", &opq_M) == 1) {
669685
vt_1 = new OPQMatrix (d, opq_M);
670-
// coarse quantizers
686+
} else if (stok == "L2norm") {
687+
vt_1 = new NormalizationTransform (d, 2.0);
688+
689+
// coarse quantizers
671690
} else if (!coarse_quantizer &&
672691
sscanf (tok, "IVF%d", &ncentroids) == 1) {
673692
if (metric == METRIC_L2) {
@@ -698,28 +717,25 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
698717
index_1 = index_ivf;
699718
} else {
700719
index_1 = new IndexFlat (d, metric);
701-
if (add_idmap) {
702-
IndexIDMap *idmap = new IndexIDMap(index_1);
703-
idmap->own_fields = true;
704-
index_1 = idmap;
705-
add_idmap = false;
706-
}
707720
}
708721
} else if (!index && (stok == "SQ8" || stok == "SQ4")) {
709-
FAISS_THROW_IF_NOT_MSG(coarse_quantizer,
710-
"ScalarQuantizer works only with an IVF");
711722
ScalarQuantizer::QuantizerType qt =
712723
stok == "SQ8" ? ScalarQuantizer::QT_8bit :
713724
stok == "SQ4" ? ScalarQuantizer::QT_4bit :
714725
ScalarQuantizer::QT_4bit;
715-
IndexIVFScalarQuantizer *index_ivf = new IndexIVFScalarQuantizer (
716-
coarse_quantizer, d, ncentroids, qt, metric);
717-
index_ivf->quantizer_trains_alone =
718-
dynamic_cast<MultiIndexQuantizer*>(coarse_quantizer)
719-
!= nullptr;
720-
del_coarse_quantizer.release ();
721-
index_ivf->own_fields = true;
722-
index_1 = index_ivf;
726+
if (coarse_quantizer) {
727+
IndexIVFScalarQuantizer *index_ivf =
728+
new IndexIVFScalarQuantizer (
729+
coarse_quantizer, d, ncentroids, qt, metric);
730+
index_ivf->quantizer_trains_alone =
731+
dynamic_cast<MultiIndexQuantizer*>(coarse_quantizer)
732+
!= nullptr;
733+
del_coarse_quantizer.release ();
734+
index_ivf->own_fields = true;
735+
index_1 = index_ivf;
736+
} else {
737+
index_1 = new IndexScalarQuantizer (d, qt, metric);
738+
}
723739
} else if (!index && sscanf (tok, "PQ%d+%d", &M, &M2) == 2) {
724740
FAISS_THROW_IF_NOT_MSG(coarse_quantizer,
725741
"PQ with + works only with an IVF");
@@ -750,13 +766,6 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
750766
IndexPQ *index_pq = new IndexPQ (d, M, 8, metric);
751767
index_pq->do_polysemous_training = true;
752768
index_1 = index_pq;
753-
if (add_idmap) {
754-
IndexIDMap *idmap = new IndexIDMap(index_1);
755-
del_index.set (idmap);
756-
idmap->own_fields = true;
757-
index_1 = idmap;
758-
add_idmap = false;
759-
}
760769
}
761770
} else if (stok == "RFlat") {
762771
make_IndexRefineFlat = true;
@@ -765,9 +774,16 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
765774
tok, description_in);
766775
}
767776

777+
if (index_1 && add_idmap) {
778+
IndexIDMap *idmap = new IndexIDMap(index_1);
779+
del_index.set (idmap);
780+
idmap->own_fields = true;
781+
index_1 = idmap;
782+
add_idmap = false;
783+
}
784+
768785
if (vt_1) {
769-
vt = vt_1;
770-
del_vt.set (vt);
786+
vts.chain.push_back (vt_1);
771787
}
772788

773789
if (coarse_quantizer_1) {
@@ -793,10 +809,14 @@ Index *index_factory (int d, const char *description_in, MetricType metric)
793809
"IDMap option not used\n");
794810
}
795811

796-
if (vt) {
797-
IndexPreTransform *index_pt = new IndexPreTransform (vt, index);
798-
del_vt.release ();
812+
if (vts.chain.size() > 0) {
813+
IndexPreTransform *index_pt = new IndexPreTransform (index);
799814
index_pt->own_fields = true;
815+
// add from back
816+
while (vts.chain.size() > 0) {
817+
index_pt->prepend_transform (vts.chain.back());
818+
vts.chain.pop_back ();
819+
}
800820
index = index_pt;
801821
}
802822

‎AuxIndexStructures.cpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,10 @@ void RangeSearchPartialResult::set_result (bool incremental)
158158
}
159159

160160

161+
/***********************************************************************
162+
* IDSelectorRange
163+
***********************************************************************/
164+
161165
IDSelectorRange::IDSelectorRange (idx_t imin, idx_t imax):
162166
imin (imin), imax (imax)
163167
{
@@ -169,6 +173,9 @@ bool IDSelectorRange::is_member (idx_t id) const
169173
}
170174

171175

176+
/***********************************************************************
177+
* IDSelectorBatch
178+
***********************************************************************/
172179

173180
IDSelectorBatch::IDSelectorBatch (long n, const idx_t *indices)
174181
{

‎AuxIndexStructures.h‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,7 @@
1515
#define FAISS_AUX_INDEX_STRUCTURES_H
1616

1717
#include <vector>
18-
19-
#if __cplusplus >= 201103L
2018
#include <unordered_set>
21-
#endif
22-
23-
#include <set>
2419

2520

2621
#include "Index.h"
@@ -80,11 +75,7 @@ struct IDSelectorRange: IDSelector {
8075
* hash collisions if lsb's are always the same */
8176
struct IDSelectorBatch: IDSelector {
8277

83-
#if __cplusplus >= 201103L
8478
std::unordered_set<idx_t> set;
85-
#else
86-
std::set<idx_t> set;
87-
#endif
8879

8980
typedef unsigned char uint8_t;
9081
std::vector<uint8_t> bloom; // assumes low bits of id are a good hash value

‎FaissException.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
// Copyright 2004-present Facebook. All Rights Reserved.
1010

1111
#include "FaissException.h"
12-
#include <cstdio>
1312

1413
namespace faiss {
1514

@@ -28,4 +27,9 @@ FaissException::FaissException(const std::string& m,
2827
funcName, file, line, m.c_str());
2928
}
3029

30+
const char*
31+
FaissException::what() const noexcept {
32+
return msg.c_str();
33+
}
34+
3135
}

‎FaissException.h‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,7 @@ class FaissException : public std::exception {
2727
int line);
2828

2929
/// from std::exception
30-
const char* what() const noexcept override
31-
{ return msg.c_str(); }
32-
~FaissException () noexcept override {}
30+
const char* what() const noexcept override;
3331

3432
std::string msg;
3533
};

‎IndexIVF.cpp‎

Lines changed: 64 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -65,21 +65,28 @@ void IndexIVF::add (idx_t n, const float * x)
6565
add_with_ids (n, x, nullptr);
6666
}
6767

68-
void IndexIVF::make_direct_map ()
68+
void IndexIVF::make_direct_map (bool new_maintain_direct_map)
6969
{
70-
if (maintain_direct_map) return;
71-
72-
direct_map.resize (ntotal, -1);
73-
for (size_t key = 0; key < nlist; key++) {
74-
const std::vector<long> & idlist = ids[key];
75-
76-
for (long ofs = 0; ofs < idlist.size(); ofs++) {
77-
direct_map [idlist [ofs]] =
78-
key << 32 | ofs;
70+
// nothing to do
71+
if (new_maintain_direct_map == maintain_direct_map)
72+
return;
73+
74+
if (new_maintain_direct_map) {
75+
direct_map.resize (ntotal, -1);
76+
for (size_t key = 0; key < nlist; key++) {
77+
const std::vector<long> & idlist = ids[key];
78+
79+
for (long ofs = 0; ofs < idlist.size(); ofs++) {
80+
FAISS_THROW_IF_NOT_MSG (
81+
0 <= idlist [ofs] && idlist[ofs] < ntotal,
82+
"direct map supported only for seuquential ids");
83+
direct_map [idlist [ofs]] = key << 32 | ofs;
84+
}
7985
}
86+
} else {
87+
direct_map.clear ();
8088
}
81-
82-
maintain_direct_map = true;
89+
maintain_direct_map = new_maintain_direct_map;
8390
}
8491

8592

@@ -183,7 +190,6 @@ void IndexIVF::merge_from (IndexIVF &other, idx_t add_id)
183190

184191

185192

186-
187193
IndexIVF::~IndexIVF()
188194
{
189195
if (own_fields) delete quantizer;
@@ -217,6 +223,8 @@ void IndexIVFFlat::add_core (idx_t n, const float * x, const long *xids,
217223

218224
{
219225
FAISS_THROW_IF_NOT (is_trained);
226+
FAISS_THROW_IF_NOT_MSG (!(maintain_direct_map && xids),
227+
"cannot have direct map and add with ids");
220228
const long * idx;
221229
ScopeDeleter<long> del;
222230

@@ -477,6 +485,49 @@ void IndexIVFFlat::copy_subset_to (IndexIVFFlat & other, int subset_type,
477485
}
478486
}
479487

488+
void IndexIVFFlat::update_vectors (int n, idx_t *new_ids, const float *x)
489+
{
490+
FAISS_THROW_IF_NOT (maintain_direct_map);
491+
FAISS_THROW_IF_NOT (is_trained);
492+
std::vector<idx_t> assign (n);
493+
quantizer->assign (n, x, assign.data());
494+
495+
for (int i = 0; i < n; i++) {
496+
idx_t id = new_ids[i];
497+
FAISS_THROW_IF_NOT_MSG (0 <= id && id < ntotal,
498+
"id to update out of range");
499+
{ // remove old one
500+
long dm = direct_map[id];
501+
long ofs = dm & 0xffffffff;
502+
long il = dm >> 32;
503+
size_t l = ids[il].size();
504+
if (ofs != l - 1) {
505+
long id2 = ids[il].back();
506+
ids[il][ofs] = id2;
507+
direct_map[id2] = (il << 32) | ofs;
508+
memcpy (vecs[il].data() + ofs * d,
509+
vecs[il].data() + (l - 1) * d,
510+
d * sizeof(vecs[il][0]));
511+
}
512+
ids[il].pop_back();
513+
vecs[il].resize((l - 1) * d);
514+
}
515+
{ // insert new one
516+
long il = assign[i];
517+
size_t l = ids[il].size();
518+
long dm = (il << 32) | l;
519+
direct_map[id] = dm;
520+
ids[il].push_back (id);
521+
vecs[il].resize((l + 1) * d);
522+
memcpy (vecs[il].data() + l * d,
523+
x + i * d,
524+
d * sizeof(vecs[il][0]));
525+
}
526+
}
527+
528+
}
529+
530+
480531

481532

482533
void IndexIVFFlat::reset()

‎IndexIVF.h‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,12 @@ struct IndexIVF: Index {
9191
size_t get_list_size (size_t list_no) const
9292
{ return ids[list_no].size(); }
9393

94-
95-
/// intialize a direct map
96-
void make_direct_map ();
94+
/** intialize a direct map
95+
*
96+
* @param new_maintain_direct_map if true, create a direct map,
97+
* else clear it
98+
*/
99+
void make_direct_map (bool new_maintain_direct_map=true);
97100

98101
/// 1= perfectly balanced, >1: imbalanced
99102
double imbalance_factor () const;
@@ -184,6 +187,16 @@ struct IndexIVFFlat: IndexIVF {
184187
const long * keys,
185188
float_maxheap_array_t * res) const;
186189

190+
/** Update a subset of vectors.
191+
*
192+
* The index must have a direct_map
193+
*
194+
* @param nv nb of vectors to update
195+
* @param idx vector indices to update, size nv
196+
* @param v vectors of new values, size nv*d
197+
*/
198+
void update_vectors (int nv, idx_t *idx, const float *v);
199+
187200
void reconstruct(idx_t key, float* recons) const override;
188201

189202
void merge_from_residuals(IndexIVF& other) override;

0 commit comments

Comments
 (0)