Skip to content

Commit d8eeceb

Browse files
mdouzemdouze
committed
sync with FB version 2017-09-08
fix index_cpu_to_gpu_multiple fix PCA dimension
1 parent 3986a8d commit d8eeceb

102 files changed

Lines changed: 7694 additions & 10565 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎IndexIVF.cpp‎

Lines changed: 174 additions & 270 deletions
Large diffs are not rendered by default.

‎IndexIVF.h‎

Lines changed: 48 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ struct IndexIVF: Index {
5454

5555
std::vector < std::vector<long> > ids; ///< Inverted lists for indexes
5656

57+
size_t code_size; ///< code size per vector in bytes
58+
std::vector < std::vector<uint8_t> > codes; // binary codes, size nlist
59+
5760
/// map for direct access to the elements. Enables reconstruct().
5861
bool maintain_direct_map;
5962
std::vector <long> direct_map;
@@ -78,13 +81,52 @@ struct IndexIVF: Index {
7881
/// does nothing by default
7982
virtual void train_residual (idx_t n, const float *x);
8083

84+
85+
/** search a set of vectors, that are pre-quantized by the IVF
86+
* quantizer. Fill in the corresponding heaps with the query
87+
* results. search() calls this.
88+
*
89+
* @param n nb of vectors to query
90+
* @param x query vectors, size nx * d
91+
* @param assign coarse quantization indices, size nx * nprobe
92+
* @param centroid_dis
93+
* distances to coarse centroids, size nx * nprobe
94+
* @param distance
95+
* output distances, size n * k
96+
* @param labels output labels, size n * k
97+
* @param store_pairs store inv list index + inv list offset
98+
* instead in upper/lower 32 bit of result,
99+
* instead of ids (used for reranking).
100+
*/
101+
virtual void search_preassigned (idx_t n, const float *x, idx_t k,
102+
const idx_t *assign,
103+
const float *centroid_dis,
104+
float *distances, idx_t *labels,
105+
bool store_pairs) const = 0;
106+
107+
/** assign the vectors, then call search_preassign */
108+
virtual void search (idx_t n, const float *x, idx_t k,
109+
float *distances, idx_t *labels) const override;
110+
111+
112+
/// Dataset manipulation functions
113+
114+
long remove_ids(const IDSelector& sel) override;
115+
81116
/** moves the entries from another dataset to self. On output,
82117
* other is empty. add_id is added to all moved ids (for
83118
* sequential ids, this would be this->ntotal */
84119
virtual void merge_from (IndexIVF &other, idx_t add_id);
85120

86-
/** implemented by sub-classes */
87-
virtual void merge_from_residuals (IndexIVF &other) = 0;
121+
/** copy a subset of the entries index to the other index
122+
*
123+
* if subset_type == 0: copies ids in [a1, a2)
124+
* if subset_type == 1: copies ids if id % a1 == a2
125+
* if subset_type == 2: copies inverted lists such that a1
126+
* elements are left before and a2 elements are after
127+
*/
128+
virtual void copy_subset_to (IndexIVF & other, int subset_type,
129+
long a1, long a2) const;
88130

89131
~IndexIVF() override;
90132

@@ -127,12 +169,9 @@ extern IndexIVFFlatStats indexIVFFlat_stats;
127169

128170
/** Inverted file with stored vectors. Here the inverted file
129171
* pre-selects the vectors to be searched, but they are not otherwise
130-
* encoded.
172+
* encoded, the code array just contains the raw float entries.
131173
*/
132174
struct IndexIVFFlat: IndexIVF {
133-
/** Inverted list of original vectors. Each list is a nl * d
134-
* matrix, where nl is the nb of vectors stored in the list. */
135-
std::vector < std::vector<float> > vecs;
136175

137176
IndexIVFFlat (
138177
Index * quantizer, size_t d, size_t nlist_,
@@ -145,48 +184,18 @@ struct IndexIVFFlat: IndexIVF {
145184
/// implemented for all IndexIVF* classes
146185
void add_with_ids(idx_t n, const float* x, const long* xids) override;
147186

148-
void search(
149-
idx_t n,
150-
const float* x,
151-
idx_t k,
152-
float* distances,
153-
idx_t* labels) const override;
154-
155-
/// perform search, without computing the assignment to the quantizer
156187
void search_preassigned (idx_t n, const float *x, idx_t k,
157188
const idx_t *assign,
158-
float *distances, idx_t *labels) const;
189+
const float *centroid_dis,
190+
float *distances, idx_t *labels,
191+
bool store_pairs) const override;
159192

160193
void range_search(
161194
idx_t n,
162195
const float* x,
163196
float radius,
164197
RangeSearchResult* result) const override;
165198

166-
/** copy a subset of the entries index to the other index
167-
*
168-
* if subset_type == 0: copies ids in [a1, a2)
169-
* if subset_type == 1: copies ids if id % a1 == a2
170-
*/
171-
void copy_subset_to (IndexIVFFlat & other, int subset_type,
172-
long a1, long a2) const;
173-
174-
void reset() override;
175-
176-
long remove_ids(const IDSelector& sel) override;
177-
178-
/// Implementation of the search for the inner product metric
179-
void search_knn_inner_product (
180-
size_t nx, const float * x,
181-
const long * keys,
182-
float_minheap_array_t * res) const;
183-
184-
/// Implementation of the search for the L2 metric
185-
void search_knn_L2sqr (
186-
size_t nx, const float * x,
187-
const long * keys,
188-
float_maxheap_array_t * res) const;
189-
190199
/** Update a subset of vectors.
191200
*
192201
* The index must have a direct_map
@@ -199,37 +208,9 @@ struct IndexIVFFlat: IndexIVF {
199208

200209
void reconstruct(idx_t key, float* recons) const override;
201210

202-
void merge_from_residuals(IndexIVF& other) override;
203-
204211
IndexIVFFlat () {}
205212
};
206213

207-
struct IndexIVFFlatIPBounds: IndexIVFFlat {
208-
209-
/// nb of dimensions of pre-filter
210-
size_t fsize;
211-
212-
/// norm of remainder (dimensions fsize:d)
213-
std::vector<std::vector<float> > part_norms;
214-
215-
IndexIVFFlatIPBounds (
216-
Index * quantizer, size_t d, size_t nlist,
217-
size_t fsize);
218-
219-
void add_core(
220-
idx_t n,
221-
const float* x,
222-
const long* xids,
223-
const long* precomputed_idx) override;
224-
225-
void search(
226-
idx_t n,
227-
const float* x,
228-
idx_t k,
229-
float* distances,
230-
idx_t* labels) const override;
231-
};
232-
233214

234215

235216
} // namespace faiss

0 commit comments

Comments
 (0)