faiss 0.1.3 → 0.1.4

data/vendor/faiss/tests/test_sliding_ivf.cpp

@@ -1,246 +0,0 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-#include <cstdio>
-#include <cstdlib>
-
-#include <memory>
-#include <random>
-#include <vector>
-
-#include <gtest/gtest.h>
-
-#include <faiss/IndexIVF.h>
-#include <faiss/AutoTune.h>
-#include <faiss/index_factory.h>
-#include <faiss/clone_index.h>
-#include <faiss/IVFlib.h>
-
-
-using namespace faiss;
-
-typedef Index::idx_t idx_t;
-
-
-// dimension of the vectors to index
-int d = 32;
-
-// nb of training vectors
-size_t nt = 5000;
-
-// size of the database points per window step
-size_t nb = 1000;
-
-// nb of queries
-size_t nq = 200;
-
-
-int total_size = 40;
-int window_size = 10;
-
-
-
-
-
-std::vector<float> make_data(size_t n)
-{
-    std::vector <float> database (n * d);
-    std::mt19937 rng;
-    std::uniform_real_distribution<> distrib;
-    for (size_t i = 0; i < n * d; i++) {
-        database[i] = distrib(rng);
-    }
-    return database;
-}
-
-std::unique_ptr<Index> make_trained_index(const char *index_type)
-{
-    auto index = std::unique_ptr<Index>(index_factory(d, index_type));
-    auto xt = make_data(nt * d);
-    index->train(nt, xt.data());
-    ParameterSpace().set_index_parameter (index.get(), "nprobe", 4);
-    return index;
-}
-
-std::vector<idx_t> search_index(Index *index, const float *xq) {
-    int k = 10;
-    std::vector<idx_t> I(k * nq);
-    std::vector<float> D(k * nq);
-    index->search (nq, xq, k, D.data(), I.data());
-    return I;
-}
-
-
-
-
-
-/*************************************************************
- * Test functions for a given index type
- *************************************************************/
-
-
-// make a few slices of indexes that can be merged
-void make_index_slices (const Index* trained_index,
-                        std::vector<std::unique_ptr<Index> > & sub_indexes) {
-
-    for (int i = 0; i < total_size; i++) {
-        sub_indexes.emplace_back (clone_index (trained_index));
-
-        printf ("preparing sub-index # %d\n", i);
-
-        Index * index = sub_indexes.back().get();
-
-        auto xb = make_data(nb * d);
-        std::vector<faiss::Index::idx_t> ids (nb);
-        std::mt19937 rng;
-        std::uniform_int_distribution<> distrib;
-        for (int j = 0; j < nb; j++) {
-            ids[j] = distrib(rng);
-        }
-        index->add_with_ids (nb, xb.data(), ids.data());
-    }
-
-}
-
-// build merged index explicitly at sliding window position i
-Index *make_merged_index(
-         const Index* trained_index,
-         const std::vector<std::unique_ptr<Index> > & sub_indexes,
-         int i) {
-
-    Index * merged_index = clone_index (trained_index);
-    for (int j = i - window_size + 1; j <= i; j++) {
-        if (j < 0 || j >= total_size) continue;
-        std::unique_ptr<Index> sub_index (
-                clone_index (sub_indexes[j].get()));
-        IndexIVF *ivf0 = ivflib::extract_index_ivf (merged_index);
-        IndexIVF *ivf1 = ivflib::extract_index_ivf (sub_index.get());
-        ivf0->merge_from (*ivf1, 0);
-        merged_index->ntotal = ivf0->ntotal;
-    }
-    return merged_index;
-}
-
-int test_sliding_window (const char *index_key) {
-
-    std::unique_ptr<Index> trained_index = make_trained_index(index_key);
-
-    // make the index slices
-    std::vector<std::unique_ptr<Index> > sub_indexes;
-
-    make_index_slices (trained_index.get(), sub_indexes);
-
-    // now slide over the windows
-    std::unique_ptr<Index> index (clone_index (trained_index.get()));
-    ivflib::SlidingIndexWindow window (index.get());
-
-    auto xq = make_data (nq * d);
-
-    for (int i = 0; i < total_size + window_size; i++) {
-
-        printf ("doing step %d / %d\n", i, total_size + window_size);
-
-        // update the index
-        window.step (i < total_size ? sub_indexes[i].get() : nullptr,
-                     i >= window_size);
-        printf ("   current n_slice = %d\n", window.n_slice);
-
-        auto new_res = search_index (index.get(), xq.data());
-
-        std::unique_ptr<Index> merged_index (
-             make_merged_index (trained_index.get(), sub_indexes, i));
-
-        auto ref_res = search_index (merged_index.get(), xq.data ());
-
-        EXPECT_EQ (ref_res.size(), new_res.size());
-
-        EXPECT_EQ (ref_res, new_res);
-    }
-    return 0;
-}
-
-
-int test_sliding_invlists (const char *index_key) {
-
-    std::unique_ptr<Index> trained_index = make_trained_index(index_key);
-
-    // make the index slices
-    std::vector<std::unique_ptr<Index> > sub_indexes;
-
-    make_index_slices (trained_index.get(), sub_indexes);
-
-    // now slide over the windows
-    std::unique_ptr<Index> index (clone_index (trained_index.get()));
-    IndexIVF * index_ivf = ivflib::extract_index_ivf (index.get());
-
-    auto xq = make_data (nq * d);
-
-    for (int i = 0; i < total_size + window_size; i++) {
-
-        printf ("doing step %d / %d\n", i, total_size + window_size);
-
-        // update the index
-        std::vector<const InvertedLists*> ils;
-        for (int j = i - window_size + 1; j <= i; j++) {
-            if (j < 0 || j >= total_size) continue;
-            ils.push_back (ivflib::extract_index_ivf (
-                      sub_indexes[j].get())->invlists);
-        }
-        if (ils.size() == 0) continue;
-
-        ConcatenatedInvertedLists *ci =
-            new ConcatenatedInvertedLists (ils.size(), ils.data());
-
-        // will be deleted by the index
-        index_ivf->replace_invlists (ci, true);
-
-        printf ("   nb invlists = %zd\n", ils.size());
-
-        auto new_res = search_index (index.get(), xq.data());
-
-        std::unique_ptr<Index> merged_index (
-             make_merged_index (trained_index.get(), sub_indexes, i));
-
-        auto ref_res = search_index (merged_index.get(), xq.data ());
-
-        EXPECT_EQ (ref_res.size(), new_res.size());
-
-        size_t ndiff = 0;
-        for (size_t j = 0; j < ref_res.size(); j++) {
-            if (ref_res[j] != new_res[j])
-                ndiff++;
-        }
-        printf("  nb differences: %zd / %zd\n",
-               ndiff, ref_res.size());
-        EXPECT_EQ (ref_res, new_res);
-    }
-    return 0;
-}
-
-
-
-
-
-/*************************************************************
- * Test entry points
- *************************************************************/
-
-TEST(SlidingWindow, IVFFlat) {
-    test_sliding_window ("IVF32,Flat");
-}
-
-TEST(SlidingWindow, PCAIVFFlat) {
-    test_sliding_window ("PCA24,IVF32,Flat");
-}
-
-TEST(SlidingInvlists, IVFFlat) {
-    test_sliding_invlists ("IVF32,Flat");
-}
-
-TEST(SlidingInvlists, PCAIVFFlat) {
-    test_sliding_invlists ("PCA24,IVF32,Flat");
-}

data/vendor/faiss/tests/test_threaded_index.cpp

@@ -1,253 +0,0 @@
-/**
- * Copyright (c) Facebook, Inc. and its 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/impl/ThreadedIndex.h>
-#include <faiss/IndexReplicas.h>
-#include <faiss/IndexShards.h>
-
-#include <chrono>
-#include <gtest/gtest.h>
-#include <memory>
-#include <vector>
-#include <thread>
-
-namespace {
-
-struct TestException : public std::exception { };
-
-struct MockIndex : public faiss::Index {
-  explicit MockIndex(idx_t d) :
-      faiss::Index(d) {
-    resetMock();
-  }
-
-  void resetMock() {
-    flag = false;
-    nCalled = 0;
-    xCalled = nullptr;
-    kCalled = 0;
-    distancesCalled = nullptr;
-    labelsCalled = nullptr;
-  }
-
-  void add(idx_t n, const float* x) override {
-    nCalled = n;
-    xCalled = x;
-  }
-
-  void search(idx_t n,
-              const float* x,
-              idx_t k,
-              float* distances,
-              idx_t* labels) const override {
-    nCalled = n;
-    xCalled = x;
-    kCalled = k;
-    distancesCalled = distances;
-    labelsCalled = labels;
-  }
-
-  void reset() override { }
-
-  bool flag;
-
-  mutable idx_t nCalled;
-  mutable const float* xCalled;
-  mutable idx_t kCalled;
-  mutable float* distancesCalled;
-  mutable idx_t* labelsCalled;
-};
-
-template <typename IndexT>
-struct MockThreadedIndex : public faiss::ThreadedIndex<IndexT> {
-  using idx_t = faiss::Index::idx_t;
-
-  explicit MockThreadedIndex(bool threaded)
-      : faiss::ThreadedIndex<IndexT>(threaded) {
-  }
-
-  void add(idx_t, const float*) override { }
-  void search(idx_t, const float*, idx_t, float*, idx_t*) const override {}
-  void reset() override {}
-};
-
-}
-
-TEST(ThreadedIndex, SingleException) {
-  std::vector<std::unique_ptr<MockIndex>> idxs;
-
-  for (int i = 0; i < 3; ++i) {
-    idxs.emplace_back(new MockIndex(1));
-  }
-
-  auto fn =
-    [](int i, MockIndex* index) {
-      if (i == 1) {
-        throw TestException();
-      } else {
-        std::this_thread::sleep_for(std::chrono::milliseconds(i * 250));
-
-        index->flag = true;
-      }
-    };
-
-  // Try with threading and without
-  for (bool threaded : {true, false}) {
-    // clear flags
-    for (auto& idx : idxs) {
-      idx->resetMock();
-    }
-
-    MockThreadedIndex<MockIndex> ti(threaded);
-    for (auto& idx : idxs) {
-      ti.addIndex(idx.get());
-    }
-
-    // The second index should throw
-    EXPECT_THROW(ti.runOnIndex(fn), TestException);
-
-    // Index 0 and 2 should have processed
-    EXPECT_TRUE(idxs[0]->flag);
-    EXPECT_TRUE(idxs[2]->flag);
-  }
-}
-
-TEST(ThreadedIndex, MultipleException) {
-  std::vector<std::unique_ptr<MockIndex>> idxs;
-
-  for (int i = 0; i < 3; ++i) {
-    idxs.emplace_back(new MockIndex(1));
-  }
-
-  auto fn =
-    [](int i, MockIndex* index) {
-      if (i < 2) {
-        throw TestException();
-      } else {
-        std::this_thread::sleep_for(std::chrono::milliseconds(i * 250));
-
-        index->flag = true;
-      }
-    };
-
-  // Try with threading and without
-  for (bool threaded : {true, false}) {
-    // clear flags
-    for (auto& idx : idxs) {
-      idx->resetMock();
-    }
-
-    MockThreadedIndex<MockIndex> ti(threaded);
-    for (auto& idx : idxs) {
-      ti.addIndex(idx.get());
-    }
-
-    // Multiple indices threw an exception that was aggregated into a
-    // FaissException
-    EXPECT_THROW(ti.runOnIndex(fn), faiss::FaissException);
-
-    // Index 2 should have processed
-    EXPECT_TRUE(idxs[2]->flag);
-  }
-}
-
-TEST(ThreadedIndex, TestReplica) {
-  int numReplicas = 5;
-  int n = 10 * numReplicas;
-  int d = 3;
-  int k = 6;
-
-  // Try with threading and without
-  for (bool threaded : {true, false}) {
-    std::vector<std::unique_ptr<MockIndex>> idxs;
-    faiss::IndexReplicas replica(d);
-
-    for (int i = 0; i < numReplicas; ++i) {
-      idxs.emplace_back(new MockIndex(d));
-      replica.addIndex(idxs.back().get());
-    }
-
-    std::vector<float> x(n * d);
-    std::vector<float> distances(n * k);
-    std::vector<faiss::Index::idx_t> labels(n * k);
-
-    replica.add(n, x.data());
-
-    for (int i = 0; i < idxs.size(); ++i) {
-      EXPECT_EQ(idxs[i]->nCalled, n);
-      EXPECT_EQ(idxs[i]->xCalled, x.data());
-    }
-
-    for (auto& idx : idxs) {
-      idx->resetMock();
-    }
-
-    replica.search(n, x.data(), k, distances.data(), labels.data());
-
-    for (int i = 0; i < idxs.size(); ++i) {
-      auto perReplica = n / idxs.size();
-
-      EXPECT_EQ(idxs[i]->nCalled, perReplica);
-      EXPECT_EQ(idxs[i]->xCalled, x.data() + i * perReplica * d);
-      EXPECT_EQ(idxs[i]->kCalled, k);
-      EXPECT_EQ(idxs[i]->distancesCalled,
-                distances.data() + (i * perReplica) * k);
-      EXPECT_EQ(idxs[i]->labelsCalled,
-                labels.data() + (i * perReplica) * k);
-    }
-  }
-}
-
-TEST(ThreadedIndex, TestShards) {
-  int numShards = 7;
-  int d = 3;
-  int n = 10 * numShards;
-  int k = 6;
-
-  // Try with threading and without
-  for (bool threaded : {true, false}) {
-    std::vector<std::unique_ptr<MockIndex>> idxs;
-    faiss::IndexShards shards(d, threaded);
-
-    for (int i = 0; i < numShards; ++i) {
-      idxs.emplace_back(new MockIndex(d));
-      shards.addIndex(idxs.back().get());
-    }
-
-    std::vector<float> x(n * d);
-    std::vector<float> distances(n * k);
-    std::vector<faiss::Index::idx_t> labels(n * k);
-
-    shards.add(n, x.data());
-
-    for (int i = 0; i < idxs.size(); ++i) {
-      auto perShard = n / idxs.size();
-
-      EXPECT_EQ(idxs[i]->nCalled, perShard);
-      EXPECT_EQ(idxs[i]->xCalled, x.data() + i * perShard * d);
-    }
-
-    for (auto& idx : idxs) {
-      idx->resetMock();
-    }
-
-    shards.search(n, x.data(), k, distances.data(), labels.data());
-
-    for (int i = 0; i < idxs.size(); ++i) {
-      auto perShard = n / idxs.size();
-
-      EXPECT_EQ(idxs[i]->nCalled, n);
-      EXPECT_EQ(idxs[i]->xCalled, x.data());
-      EXPECT_EQ(idxs[i]->kCalled, k);
-      // There is a temporary buffer used for shards
-      EXPECT_EQ(idxs[i]->distancesCalled,
-                idxs[0]->distancesCalled + i * k * n);
-      EXPECT_EQ(idxs[i]->labelsCalled,
-                idxs[0]->labelsCalled + i * k * n);
-    }
-  }
-}