Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
test_transfer_invlists.cpp
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <cstdio>
10 #include <cstdlib>
11
12 #include <gtest/gtest.h>
13
14 #include <faiss/IndexIVFFlat.h>
15 #include <faiss/index_io.h>
16 #include <faiss/AuxIndexStructures.h>
17 #include <faiss/AutoTune.h>
18 #include <faiss/VectorTransform.h>
19 #include <faiss/utils.h>
20 #include <faiss/IVFlib.h>
21
22
23 // parameters to use for the test
24 int d = 64;
25 size_t nb = 1000;
26 size_t nq = 100;
27 size_t nt = 500;
28 int k = 10;
29 int nlist = 40;
30
31 using namespace faiss;
32
33 typedef faiss::Index::idx_t idx_t;
34
35
36 std::vector<float> get_data (size_t nb, int seed) {
37 std::vector<float> x (nb * d);
38 float_randn (x.data(), nb * d, seed);
39 return x;
40 }
41
42
43
44
45
46 void test_index_type(const char *factory_string) {
47
48 // transfer inverted lists in nslice slices
49 int nslice = 3;
50
51 /****************************************************************
52 * trained reference index
53 ****************************************************************/
54
55 std::unique_ptr<Index> trained (index_factory (d, factory_string));
56
57 {
58 auto xt = get_data (nt, 123);
59 trained->train (nt, xt.data());
60 }
61
62 // sample nq query vectors to check if results are the same
63 auto xq = get_data (nq, 818);
64
65
66 /****************************************************************
67 * source index
68 ***************************************************************/
69 std::unique_ptr<Index> src_index (clone_index (trained.get()));
70
71 { // add some data to source index
72 auto xb = get_data (nb, 245);
73 src_index->add (nb, xb.data());
74 }
75
76 ParameterSpace().set_index_parameter (src_index.get(), "nprobe", 4);
77
78 // remember reference search result on source index
79 std::vector<idx_t> Iref (nq * k);
80 std::vector<float> Dref (nq * k);
81 src_index->search (nq, xq.data(), k, Dref.data(), Iref.data());
82
83 /****************************************************************
84 * destination index -- should be replaced by source index
85 ***************************************************************/
86
87 std::unique_ptr<Index> dst_index (clone_index (trained.get()));
88
89 { // initial state: filled in with some garbage
90 int nb2 = nb + 10;
91 auto xb = get_data (nb2, 366);
92 dst_index->add (nb2, xb.data());
93 }
94
95 std::vector<idx_t> Inew (nq * k);
96 std::vector<float> Dnew (nq * k);
97
98 ParameterSpace().set_index_parameter (dst_index.get(), "nprobe", 4);
99
100 // transfer from source to destination in nslice slices
101 for (int sl = 0; sl < nslice; sl++) {
102
103 // so far, the indexes are different
104 dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
105 EXPECT_TRUE (Iref != Inew);
106 EXPECT_TRUE (Dref != Dnew);
107
108 // range of inverted list indices to transfer
109 long i0 = sl * nlist / nslice;
110 long i1 = (sl + 1) * nlist / nslice;
111
112 std::vector<uint8_t> data_to_transfer;
113 {
114 std::unique_ptr<ArrayInvertedLists> il
115 (ivflib::get_invlist_range (src_index.get(), i0, i1));
116 // serialize inverted lists
117 VectorIOWriter wr;
118 write_InvertedLists (il.get(), &wr);
119 data_to_transfer.swap (wr.data);
120 }
121
122 // transfer data here from source machine to dest machine
123
124 {
125 VectorIOReader reader;
126 reader.data.swap (data_to_transfer);
127
128 // deserialize inverted lists
129 std::unique_ptr<ArrayInvertedLists> il
130 (dynamic_cast<ArrayInvertedLists *>
131 (read_InvertedLists (&reader)));
132
133 // swap inverted lists. Block searches here!
134 {
135 ivflib::set_invlist_range (dst_index.get(), i0, i1, il.get());
136 }
137 }
138
139 }
140 EXPECT_EQ (dst_index->ntotal, src_index->ntotal);
141
142 // now, the indexes are the same
143 dst_index->search (nq, xq.data(), k, Dnew.data(), Inew.data());
144 EXPECT_TRUE (Iref == Inew);
145 EXPECT_TRUE (Dref == Dnew);
146
147 }
148
149
150 TEST(TRANS, IVFFlat) {
151 test_index_type ("IVF40,Flat");
152 }
153
154 TEST(TRANS, IVFFlatPreproc) {
155 test_index_type ("PCAR32,IVF40,Flat");
156 }
faiss::VectorIOReader
Definition: AuxIndexStructures.h:206
faiss::VectorIOWriter
Definition: AuxIndexStructures.h:212
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::ParameterSpace::set_index_parameter
virtual void set_index_parameter(Index *index, const std::string &name, double val) const
set one of the parameters
Definition: AutoTune.cpp:452
faiss::ParameterSpace
Definition: AutoTune.h:132
faiss::index_factory
Index * index_factory(int d, const char *description_in, MetricType metric)
Definition: AutoTune.cpp:722