Skip to content

Commit

Permalink
Access graph structure for NSG (facebookresearch#2984)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: facebookresearch#2984

It is not entirely trivial to access the NSG graph structure from Python (although it is a fixed size N-by-K matrix of vector ids).
This diff adds an inspect_tools function to do that.

Reviewed By: algoriddle

Differential Revision: D48026775

fbshipit-source-id: 94cd7be7f656bcd333d62586531f287ea8e052e5
  • Loading branch information
mdouze authored and facebook-github-bot committed Aug 4, 2023
1 parent da16d9d commit 687457b
Show file tree
Hide file tree
Showing 6 changed files with 59 additions and 2 deletions.
13 changes: 13 additions & 0 deletions contrib/inspect_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,3 +96,16 @@ def get_flat_data(index):
""" copy and return the data matrix in an IndexFlat """
xb = faiss.vector_to_array(index.codes).view("float32")
return xb.reshape(index.ntotal, index.d)


def get_NSG_neighbors(nsg):
""" get the neighbor list for the vectors stored in the NSG structure, as
a N-by-K matrix of indices """
graph = nsg.get_final_graph()
neighbors = np.zeros((graph.N, graph.K), dtype='int32')
faiss.memcpy(
faiss.swig_ptr(neighbors),
graph.data,
neighbors.nbytes
)
return neighbors
6 changes: 5 additions & 1 deletion faiss/impl/NNDescent.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -374,6 +374,10 @@ void NNDescent::init_graph(DistanceComputer& qdis) {

void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
FAISS_THROW_IF_NOT_MSG(L >= K, "L should be >= K in NNDescent.build");
FAISS_THROW_IF_NOT_FMT(
n > NUM_EVAL_POINTS,
"NNDescent.build cannot build a graph smaller than %d",
int(NUM_EVAL_POINTS));

if (verbose) {
printf("Parameters: K=%d, S=%d, R=%d, L=%d, iter=%d\n",
Expand Down Expand Up @@ -403,7 +407,7 @@ void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
has_built = true;

if (verbose) {
printf("Addes %d points into the index\n", ntotal);
printf("Added %d points into the index\n", ntotal);
}
}

Expand Down
2 changes: 1 addition & 1 deletion faiss/impl/NSG.h
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ namespace nsg {

template <class node_t>
struct Graph {
node_t* data; ///< the flattened adjacency matrix
node_t* data; ///< the flattened adjacency matrix, size N-by-K
int K; ///< nb of neighbors per node
int N; ///< total nb of nodes
bool own_fields; ///< the underlying data owned by itself or not
Expand Down
13 changes: 13 additions & 0 deletions faiss/python/swigfaiss.swig
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,20 @@ void gpu_sync_all_devices()
%include <faiss/IndexNNDescent.h>

%include <faiss/IndexIVFFlat.h>

%warnfilter(509) faiss::nsg::Graph< int >::at(int,int);

%include <faiss/impl/NSG.h>

%template(NSG_Graph_int) faiss::nsg::Graph<int>;

// not using %shared_ptr to avoid mem leaks
%extend faiss::NSG {
faiss::nsg::Graph<int>* get_final_graph() {
return $self->final_graph.get();
}
}

%include <faiss/IndexNSG.h>

#ifndef SWIGWIN
Expand Down
17 changes: 17 additions & 0 deletions tests/test_build_blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -542,6 +542,23 @@ def subtest(self, d, K, metric):
print('Metric: {}, knng accuracy: {}'.format(metric_names[metric], recall))
assert recall > 0.99

def test_small_nndescent(self):
""" building a too small graph used to crash, make sure it raises
an exception instead.
TODO: build the exact knn graph for small cases
"""
d = 32
K = 10
index = faiss.IndexNNDescentFlat(d, K, faiss.METRIC_L2)
index.nndescent.S = 10
index.nndescent.R = 32
index.nndescent.L = K + 20
index.nndescent.iter = 5
index.verbose = True

xb = np.zeros((78, d), dtype='float32')
self.assertRaises(RuntimeError, index.add, xb)


class TestResultHeap(unittest.TestCase):

Expand Down
10 changes: 10 additions & 0 deletions tests/test_contrib.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,16 @@ def test_make_LT(self):
Ynew = lt.apply(X)
np.testing.assert_equal(Yref, Ynew)

def test_NSG_neighbors(self):
# FIXME number of elements to add should be >> 100
ds = datasets.SyntheticDataset(32, 0, 200, 10)
index = faiss.index_factory(ds.d, "NSG")
index.add(ds.get_database())
neighbors = inspect_tools.get_NSG_neighbors(index.nsg)
# neighbors should be either valid indexes or -1
np.testing.assert_array_less(-2, neighbors)
np.testing.assert_array_less(neighbors, ds.nb)


class TestRangeEval(unittest.TestCase):

Expand Down

0 comments on commit 687457b

Please sign in to comment.