1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
|
/*
* Copyright (c) Meta Platforms, Inc. and 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/IndexNSG.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/random.h>
#include <gtest/gtest.h>
using namespace faiss;
using FinalNSGGraph = nsg::Graph<int32_t>;
struct CompressedNSGGraph : FinalNSGGraph {
int bits;
size_t stride;
std::vector<uint8_t> compressed_data;
CompressedNSGGraph(const FinalNSGGraph& graph, int bits)
: FinalNSGGraph(graph.data, graph.N, graph.K), bits(bits) {
FAISS_THROW_IF_NOT((1 << bits) >= K + 1);
stride = (K * bits + 7) / 8;
compressed_data.resize(N * stride);
for (size_t i = 0; i < N; i++) {
BitstringWriter writer(compressed_data.data() + i * stride, stride);
for (size_t j = 0; j < K; j++) {
int32_t v = graph.data[i * K + j];
if (v == -1) {
writer.write(K + 1, bits);
break;
} else {
writer.write(v, bits);
}
}
}
data = nullptr;
}
size_t get_neighbors(int i, int32_t* neighbors) const override {
BitstringReader reader(compressed_data.data() + i * stride, stride);
for (int j = 0; j < K; j++) {
int32_t v = reader.read(bits);
if (v == K + 1) {
return j;
}
neighbors[j] = v;
}
return K;
}
};
TEST(NSGCompressed, test_compressed) {
size_t nq = 10, nt = 0, nb = 5000, d = 32, k = 10;
using idx_t = faiss::idx_t;
std::vector<float> buf((nq + nb + nt) * d);
faiss::rand_smooth_vectors(nq + nb + nt, d, buf.data(), 1234);
const float* xt = buf.data();
const float* xb = xt + nt * d;
const float* xq = xb + nb * d;
faiss::IndexNSGFlat index(d, 32);
index.add(nb, xb);
std::vector<faiss::idx_t> Iref(nq * k);
std::vector<float> Dref(nq * k);
index.search(nq, xq, k, Dref.data(), Iref.data());
// replace the shared ptr
index.nsg.final_graph.reset(
new CompressedNSGGraph(*index.nsg.final_graph, 13));
std::vector<idx_t> I(nq * k);
std::vector<float> D(nq * k);
index.search(nq, xq, k, D.data(), I.data());
// make sure we find back the original results
EXPECT_EQ(Iref, I);
EXPECT_EQ(Dref, D);
}
|
