/** * Copyright (c) 2015-present, Facebook, Inc. * All rights reserved. * * This source code is licensed under the BSD+Patents license found in the * LICENSE file in the root directory of this source tree. */ #include <cstdio> #include <cstdlib> #include <cassert> #include <faiss/IndexFlat.h> #include <faiss/IndexIVF.h> int main() { int d = 64; // dimension int nb = 100000; // database size int nq = 10000; // nb of queries float *xb = new float[d * nb]; float *xq = new float[d * nq]; for(int i = 0; i < nb; i++) { for(int j = 0; j < d; j++) xb[d * i + j] = drand48(); xb[d * i] += i / 1000.; } for(int i = 0; i < nq; i++) { for(int j = 0; j < d; j++) xq[d * i + j] = drand48(); xq[d * i] += i / 1000.; } int nlist = 100; int k = 4; faiss::IndexFlatL2 quantizer(d); // the other index faiss::IndexIVFFlat index(&quantizer, d, nlist, faiss::METRIC_L2); // here we specify METRIC_L2, by default it performs inner-product search assert(!index.is_trained); index.train(nb, xb); assert(index.is_trained); index.add(nb, xb); { // search xq long *I = new long[k * nq]; float *D = new float[k * nq]; index.search(nq, xq, k, D, I); printf("I=\n"); for(int i = nq - 5; i < nq; i++) { for(int j = 0; j < k; j++) printf("%5ld ", I[i * k + j]); printf("\n"); } index.nprobe = 10; index.search(nq, xq, k, D, I); printf("I=\n"); for(int i = nq - 5; i < nq; i++) { for(int j = 0; j < k; j++) printf("%5ld ", I[i * k + j]); printf("\n"); } delete [] I; delete [] D; } delete [] xb; delete [] xq; return 0; }