Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVFPQ.h
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved.
11
12 #pragma once
13
14 #include "GpuIndexIVF.h"
15 #include <vector>
16
17 namespace faiss { struct IndexIVFPQ; }
18
19 namespace faiss { namespace gpu {
20
21 class GpuIndexFlat;
22 class IVFPQ;
23
24 /// IVFPQ index for the GPU
25 class GpuIndexIVFPQ : public GpuIndexIVF {
26 public:
27 /// Construct from a pre-existing faiss::IndexIVFPQ instance, copying
28 /// data over to the given GPU, if the input index is trained.
29 GpuIndexIVFPQ(GpuResources* resources,
30 int device,
31 IndicesOptions indicesOptions,
32 bool useFloat16LookupTables,
33 const faiss::IndexIVFPQ* index);
34
35 /// Construct an empty index
36 GpuIndexIVFPQ(GpuResources* resources,
37 int device,
38 int dims,
39 int nlist,
40 int subQuantizers,
41 int bitsPerCode,
42 bool usePrecomputed,
43 IndicesOptions indicesOptions,
44 bool useFloat16LookupTables,
45 faiss::MetricType metric);
46
47 virtual ~GpuIndexIVFPQ();
48
49 /// Reserve space on the GPU for the inverted lists for `num`
50 /// vectors, assumed equally distributed among
51
52 /// Initialize ourselves from the given CPU index; will overwrite
53 /// all data in ourselves
54 void copyFrom(const faiss::IndexIVFPQ* index);
55
56 /// Copy ourselves to the given CPU index; will overwrite all data
57 /// in the index instance
58 void copyTo(faiss::IndexIVFPQ* index) const;
59
60 /// Reserve GPU memory in our inverted lists for this number of vectors
61 void reserveMemory(size_t numVecs);
62
63 /// Enable or disable pre-computed codes
64 void setPrecomputedCodes(bool enable);
65
66 /// Are pre-computed codes enabled?
67 bool getPrecomputedCodes() const;
68
69 /// Are float16 residual distance lookup tables enabled?
70 bool getFloat16LookupTables() const;
71
72 /// Return the number of sub-quantizers we are using
73 int getNumSubQuantizers() const;
74
75 /// Return the number of bits per PQ code
76 int getBitsPerCode() const;
77
78 /// Return the number of centroids per PQ code (2^bits per code)
79 int getCentroidsPerSubQuantizer() const;
80
81 /// After adding vectors, one can call this to reclaim device memory
82 /// to exactly the amount needed. Returns space reclaimed in bytes
83 size_t reclaimMemory();
84
85 /// Clears out all inverted lists, but retains the coarse and
86 /// product centroid information
87 virtual void reset();
88
89 virtual void train(Index::idx_t n, const float* x);
90
91 virtual void set_typename();
92
93 /// For debugging purposes, return the list length of a particular
94 /// list
95 int getListLength(int listId) const;
96
97 /// For debugging purposes, return the list codes of a particular
98 /// list
99 std::vector<unsigned char> getListCodes(int listId) const;
100
101 /// For debugging purposes, return the list indices of a particular
102 /// list
103 std::vector<long> getListIndices(int listId) const;
104
105 protected:
106 /// Called from GpuIndex for add/add_with_ids
107 virtual void addImpl_(faiss::Index::idx_t n,
108 const float* x,
109 const faiss::Index::idx_t* ids);
110
111 /// Called from GpuIndex for search
112 virtual void searchImpl_(faiss::Index::idx_t n,
113 const float* x,
114 faiss::Index::idx_t k,
115 float* distances,
116 faiss::Index::idx_t* labels) const;
117
118 private:
119 void assertSettings_() const;
120
121 void trainResidualQuantizer_(Index::idx_t n, const float* x);
122
123 private:
124 /// Do we use float16 residual distance lookup tables for query?
125 const bool useFloat16LookupTables_;
126
127 /// Number of sub-quantizers per encoded vector
128 int subQuantizers_;
129
130 /// Bits per sub-quantizer code
131 int bitsPerCode_;
132
133 /// Should we or should we not use precomputed codes?
134 bool usePrecomputed_;
135
136 /// Desired inverted list memory reservation
137 size_t reserveMemoryVecs_;
138
139 /// The product quantizer instance that we own; contains the
140 /// inverted lists
141 IVFPQ* index_;
142 };
143
144 } } // namespace
faiss::gpu::GpuIndexIVFPQ::getListIndices
std::vector< long > getListIndices(int listId) const
Definition: GpuIndexIVFPQ.cu:421
faiss::gpu::GpuIndexIVFPQ::getListLength
int getListLength(int listId) const
Definition: GpuIndexIVFPQ.cu:407
faiss::gpu::GpuIndexIVFPQ::addImpl_
virtual void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids)
Called from GpuIndex for add/add_with_ids.
Definition: GpuIndexIVFPQ.cu:326
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::GpuIndexIVFPQ::getBitsPerCode
int getBitsPerCode() const
Return the number of bits per PQ code.
Definition: GpuIndexIVFPQ.cu:232
faiss::gpu::GpuIndexIVF
Definition: GpuIndexIVF.h:25
faiss::gpu::GpuIndexIVFPQ::reclaimMemory
size_t reclaimMemory()
Definition: GpuIndexIVFPQ.cu:242
faiss::gpu::GpuIndexIVFPQ::searchImpl_
virtual void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const
Called from GpuIndex for search.
Definition: GpuIndexIVFPQ.cu:356
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:30
faiss::gpu::GpuIndexIVFPQ::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: GpuIndexIVFPQ.cu:197
faiss::gpu::GpuIndexIVFPQ::reset
virtual void reset()
Definition: GpuIndexIVFPQ.cu:252
faiss::gpu::GpuIndexIVFPQ::copyFrom
void copyFrom(const faiss::IndexIVFPQ *index)
Definition: GpuIndexIVFPQ.cu:92
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::gpu::GpuIndexIVFPQ::copyTo
void copyTo(faiss::IndexIVFPQ *index) const
Definition: GpuIndexIVFPQ.cu:149
faiss::gpu::GpuIndexIVFPQ::getFloat16LookupTables
bool getFloat16LookupTables() const
Are float16 residual distance lookup tables enabled?
Definition: GpuIndexIVFPQ.cu:222
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:25
faiss::gpu::GpuIndexIVFPQ::getNumSubQuantizers
int getNumSubQuantizers() const
Return the number of sub-quantizers we are using.
Definition: GpuIndexIVFPQ.cu:227
faiss::gpu::GpuIndexIVFPQ::getCentroidsPerSubQuantizer
int getCentroidsPerSubQuantizer() const
Return the number of centroids per PQ code (2^bits per code)
Definition: GpuIndexIVFPQ.cu:237
faiss::gpu::GpuIndexIVFPQ::setPrecomputedCodes
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
Definition: GpuIndexIVFPQ.cu:206
faiss::gpu::GpuIndexIVFPQ::getPrecomputedCodes
bool getPrecomputedCodes() const
Are pre-computed codes enabled?
Definition: GpuIndexIVFPQ.cu:217
faiss::gpu::GpuIndexIVFPQ::GpuIndexIVFPQ
GpuIndexIVFPQ(GpuResources *resources, int device, IndicesOptions indicesOptions, bool useFloat16LookupTables, const faiss::IndexIVFPQ *index)
Definition: GpuIndexIVFPQ.cu:26
faiss::gpu::IVFPQ
Implementing class for IVFPQ on the GPU.
Definition: IVFPQ.cuh:20
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44
faiss::gpu::GpuIndexIVFPQ::train
virtual void train(Index::idx_t n, const float *x)
Definition: GpuIndexIVFPQ.cu:307
faiss::gpu::GpuIndexIVFPQ::getListCodes
std::vector< unsigned char > getListCodes(int listId) const
Definition: GpuIndexIVFPQ.cu:413