faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
-
 #include <cuda_profiler_api.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFPQ.h>
@@ -32,108 +30,131 @@ DEFINE_bool(per_batch_time, false, "print per-batch times");
 DEFINE_bool(reserve_memory, false, "whether or not to pre-reserve memory");
 
 int main(int argc, char** argv) {
-  gflags::ParseCommandLineFlags(&argc, &argv, true);
-
-  cudaProfilerStop();
-
-  int dim = FLAGS_dim;
-  int numCentroids = FLAGS_centroids;
-  int bytesPerVec = FLAGS_bytes_per_vec;
-  int bitsPerCode = FLAGS_bits_per_code;
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-  faiss::gpu::StandardGpuResources res;
+    cudaProfilerStop();
 
-  // IndexIVFPQ will complain, but just give us enough to get through this
-  int numTrain = 4 * numCentroids;
-  std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
+    int dim = FLAGS_dim;
+    int numCentroids = FLAGS_centroids;
+    int bytesPerVec = FLAGS_bytes_per_vec;
+    int bitsPerCode = FLAGS_bits_per_code;
 
-  faiss::IndexFlatL2 coarseQuantizer(dim);
-  faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, dim, numCentroids,
-                             bytesPerVec, bitsPerCode);
-  if (FLAGS_time_cpu) {
-    cpuIndex.train(numTrain, trainVecs.data());
-  }
+    faiss::gpu::StandardGpuResources res;
 
-  faiss::gpu::GpuIndexIVFPQConfig config;
-  config.device = 0;
-  config.indicesOptions = (faiss::gpu::IndicesOptions) FLAGS_index;
+    // IndexIVFPQ will complain, but just give us enough to get through this
+    int numTrain = 4 * numCentroids;
+    std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
 
-  faiss::gpu::GpuIndexIVFPQ gpuIndex(
-    &res, dim, numCentroids, bytesPerVec, bitsPerCode,
-    faiss::METRIC_L2, config);
-
-  if (FLAGS_time_gpu) {
-    gpuIndex.train(numTrain, trainVecs.data());
-    if (FLAGS_reserve_memory) {
-      size_t numVecs = (size_t) FLAGS_batches * (size_t) FLAGS_batch_size;
-      gpuIndex.reserveMemory(numVecs);
+    faiss::IndexFlatL2 coarseQuantizer(dim);
+    faiss::IndexIVFPQ cpuIndex(
+            &coarseQuantizer, dim, numCentroids, bytesPerVec, bitsPerCode);
+    if (FLAGS_time_cpu) {
+        cpuIndex.train(numTrain, trainVecs.data());
     }
-  }
 
-  cudaDeviceSynchronize();
-  CUDA_VERIFY(cudaProfilerStart());
+    faiss::gpu::GpuIndexIVFPQConfig config;
+    config.device = 0;
+    config.indicesOptions = (faiss::gpu::IndicesOptions)FLAGS_index;
 
-  float totalGpuTime = 0.0f;
-  float totalCpuTime = 0.0f;
+    faiss::gpu::GpuIndexIVFPQ gpuIndex(
+            &res,
+            dim,
+            numCentroids,
+            bytesPerVec,
+            bitsPerCode,
+            faiss::METRIC_L2,
+            config);
 
-  for (int i = 0; i < FLAGS_batches; ++i) {
-    if (!FLAGS_per_batch_time) {
-      if (i % 10 == 0) {
-        printf("Adding batch %d\n", i + 1);
-      }
+    if (FLAGS_time_gpu) {
+        gpuIndex.train(numTrain, trainVecs.data());
+        if (FLAGS_reserve_memory) {
+            size_t numVecs = (size_t)FLAGS_batches * (size_t)FLAGS_batch_size;
+            gpuIndex.reserveMemory(numVecs);
+        }
     }
 
-    auto addVecs = faiss::gpu::randVecs(FLAGS_batch_size, dim);
+    cudaDeviceSynchronize();
+    CUDA_VERIFY(cudaProfilerStart());
+
+    float totalGpuTime = 0.0f;
+    float totalCpuTime = 0.0f;
+
+    for (int i = 0; i < FLAGS_batches; ++i) {
+        if (!FLAGS_per_batch_time) {
+            if (i % 10 == 0) {
+                printf("Adding batch %d\n", i + 1);
+            }
+        }
+
+        auto addVecs = faiss::gpu::randVecs(FLAGS_batch_size, dim);
+
+        if (FLAGS_time_gpu) {
+            faiss::gpu::CpuTimer timer;
+            gpuIndex.add(FLAGS_batch_size, addVecs.data());
+            CUDA_VERIFY(cudaDeviceSynchronize());
+            auto time = timer.elapsedMilliseconds();
+
+            totalGpuTime += time;
+
+            if (FLAGS_per_batch_time) {
+                printf("Batch %d | GPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
+                       i + 1,
+                       FLAGS_batch_size,
+                       time,
+                       time / (float)FLAGS_batch_size);
+            }
+        }
+
+        if (FLAGS_time_cpu) {
+            faiss::gpu::CpuTimer timer;
+            cpuIndex.add(FLAGS_batch_size, addVecs.data());
+            auto time = timer.elapsedMilliseconds();
+
+            totalCpuTime += time;
+
+            if (FLAGS_per_batch_time) {
+                printf("Batch %d | CPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
+                       i + 1,
+                       FLAGS_batch_size,
+                       time,
+                       time / (float)FLAGS_batch_size);
+            }
+        }
+    }
 
-    if (FLAGS_time_gpu) {
-      faiss::gpu::CpuTimer timer;
-      gpuIndex.add(FLAGS_batch_size, addVecs.data());
-      CUDA_VERIFY(cudaDeviceSynchronize());
-      auto time = timer.elapsedMilliseconds();
+    CUDA_VERIFY(cudaProfilerStop());
 
-      totalGpuTime += time;
+    int total = FLAGS_batch_size * FLAGS_batches;
 
-      if (FLAGS_per_batch_time) {
-      printf("Batch %d | GPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
-             i + 1, FLAGS_batch_size, time, time / (float) FLAGS_batch_size);
-      }
+    if (FLAGS_time_gpu) {
+        printf("%d dim, %d centroids, %d x %d encoding\n"
+               "GPU time to add %d vectors (%d batches, %d per batch): "
+               "%.3f ms (%.3f us per)\n",
+               dim,
+               numCentroids,
+               bytesPerVec,
+               bitsPerCode,
+               total,
+               FLAGS_batches,
+               FLAGS_batch_size,
+               totalGpuTime,
+               totalGpuTime * 1000.0f / (float)total);
     }
 
     if (FLAGS_time_cpu) {
-      faiss::gpu::CpuTimer timer;
-      cpuIndex.add(FLAGS_batch_size, addVecs.data());
-      auto time = timer.elapsedMilliseconds();
-
-      totalCpuTime += time;
-
-      if (FLAGS_per_batch_time) {
-        printf("Batch %d | CPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
-               i + 1, FLAGS_batch_size, time, time / (float) FLAGS_batch_size);
-      }
+        printf("%d dim, %d centroids, %d x %d encoding\n"
+               "CPU time to add %d vectors (%d batches, %d per batch): "
+               "%.3f ms (%.3f us per)\n",
+               dim,
+               numCentroids,
+               bytesPerVec,
+               bitsPerCode,
+               total,
+               FLAGS_batches,
+               FLAGS_batch_size,
+               totalCpuTime,
+               totalCpuTime * 1000.0f / (float)total);
     }
-  }
-
-  CUDA_VERIFY(cudaProfilerStop());
-
-  int total = FLAGS_batch_size * FLAGS_batches;
-
-  if (FLAGS_time_gpu) {
-    printf("%d dim, %d centroids, %d x %d encoding\n"
-           "GPU time to add %d vectors (%d batches, %d per batch): "
-           "%.3f ms (%.3f us per)\n",
-           dim, numCentroids, bytesPerVec, bitsPerCode,
-           total, FLAGS_batches, FLAGS_batch_size,
-           totalGpuTime, totalGpuTime * 1000.0f / (float) total);
-  }
-
-  if (FLAGS_time_cpu) {
-    printf("%d dim, %d centroids, %d x %d encoding\n"
-           "CPU time to add %d vectors (%d batches, %d per batch): "
-           "%.3f ms (%.3f us per)\n",
-           dim, numCentroids, bytesPerVec, bitsPerCode,
-           total, FLAGS_batches, FLAGS_batch_size,
-           totalCpuTime, totalCpuTime * 1000.0f / (float) total);
-  }
-
-  return 0;
+
+    return 0;
 }

data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp

@@ -5,14 +5,13 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
+#include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
-#include <faiss/IndexFlat.h>
-#include <faiss/index_io.h>
 #include <faiss/gpu/test/TestUtils.h>
-#include <vector>
+#include <faiss/index_io.h>
 #include <gflags/gflags.h>
+#include <vector>
 
 // For IVFPQ:
 DEFINE_bool(ivfpq, false, "use IVFPQ encoding");
@@ -32,71 +31,83 @@ DEFINE_int32(num_train, -1, "number of database vecs to train on");
 
 template <typename T>
 void fillAndSave(T& index, int numTrain, int num, int dim) {
-  auto trainVecs = faiss::gpu::randVecs(numTrain, dim);
-  index.train(numTrain, trainVecs.data());
+    auto trainVecs = faiss::gpu::randVecs(numTrain, dim);
+    index.train(numTrain, trainVecs.data());
 
-  constexpr int kAddChunk = 1000000;
+    constexpr int kAddChunk = 1000000;
 
-  for (int i = 0; i < num; i += kAddChunk) {
-    int numRemaining = (num - i) < kAddChunk ? (num - i) : kAddChunk;
-    auto vecs = faiss::gpu::randVecs(numRemaining, dim);
+    for (int i = 0; i < num; i += kAddChunk) {
+        int numRemaining = (num - i) < kAddChunk ? (num - i) : kAddChunk;
+        auto vecs = faiss::gpu::randVecs(numRemaining, dim);
 
-    printf("adding at %d: %d\n", i, numRemaining);
-    index.add(numRemaining, vecs.data());
-  }
+        printf("adding at %d: %d\n", i, numRemaining);
+        index.add(numRemaining, vecs.data());
+    }
 
-  faiss::write_index(&index, FLAGS_out.c_str());
+    faiss::write_index(&index, FLAGS_out.c_str());
 }
 
 int main(int argc, char** argv) {
-  gflags::ParseCommandLineFlags(&argc, &argv, true);
-
-  // Either ivfpq or ivfflat must be set
-  if ((FLAGS_ivfpq && FLAGS_ivfflat) ||
-      (!FLAGS_ivfpq && !FLAGS_ivfflat)) {
-    printf("must specify either ivfpq or ivfflat\n");
-    return 1;
-  }
-
-  auto dim = FLAGS_dim;
-  auto numCentroids = FLAGS_num_coarse;
-  auto num = FLAGS_num;
-  auto numTrain = FLAGS_num_train;
-  numTrain = numTrain == -1 ? std::max((num / 4), 1) : numTrain;
-  numTrain = std::min(num, numTrain);
-
-  if (FLAGS_ivfpq) {
-    faiss::IndexFlatL2 quantizer(dim);
-    faiss::IndexIVFPQ index(&quantizer, dim, numCentroids,
-                            FLAGS_codes, FLAGS_bits_per_code);
-    index.verbose = true;
-
-    printf("IVFPQ: codes %d bits per code %d\n",
-           FLAGS_codes, FLAGS_bits_per_code);
-    printf("Lists: %d\n", numCentroids);
-    printf("Database: dim %d num vecs %d trained on %d\n", dim, num, numTrain);
-    printf("output file: %s\n", FLAGS_out.c_str());
-
-    fillAndSave(index, numTrain, num, dim);
-  } else if (FLAGS_ivfflat) {
-    faiss::IndexFlatL2 quantizerL2(dim);
-    faiss::IndexFlatIP quantizerIP(dim);
-
-    faiss::IndexFlat* quantizer = FLAGS_l2 ?
-      (faiss::IndexFlat*) &quantizerL2 :
-      (faiss::IndexFlat*) &quantizerIP;
-
-    faiss::IndexIVFFlat index(quantizer, dim, numCentroids,
-                              FLAGS_l2 ? faiss::METRIC_L2 :
-                              faiss::METRIC_INNER_PRODUCT);
-
-    printf("IVFFlat: metric %s\n", FLAGS_l2 ? "L2" : "IP");
-    printf("Lists: %d\n", numCentroids);
-    printf("Database: dim %d num vecs %d trained on %d\n", dim, num, numTrain);
-    printf("output file: %s\n", FLAGS_out.c_str());
-
-    fillAndSave(index, numTrain, num, dim);
-  }
-
-  return 0;
+    gflags::ParseCommandLineFlags(&argc, &argv, true);
+
+    // Either ivfpq or ivfflat must be set
+    if ((FLAGS_ivfpq && FLAGS_ivfflat) || (!FLAGS_ivfpq && !FLAGS_ivfflat)) {
+        printf("must specify either ivfpq or ivfflat\n");
+        return 1;
+    }
+
+    auto dim = FLAGS_dim;
+    auto numCentroids = FLAGS_num_coarse;
+    auto num = FLAGS_num;
+    auto numTrain = FLAGS_num_train;
+    numTrain = numTrain == -1 ? std::max((num / 4), 1) : numTrain;
+    numTrain = std::min(num, numTrain);
+
+    if (FLAGS_ivfpq) {
+        faiss::IndexFlatL2 quantizer(dim);
+        faiss::IndexIVFPQ index(
+                &quantizer,
+                dim,
+                numCentroids,
+                FLAGS_codes,
+                FLAGS_bits_per_code);
+        index.verbose = true;
+
+        printf("IVFPQ: codes %d bits per code %d\n",
+               FLAGS_codes,
+               FLAGS_bits_per_code);
+        printf("Lists: %d\n", numCentroids);
+        printf("Database: dim %d num vecs %d trained on %d\n",
+               dim,
+               num,
+               numTrain);
+        printf("output file: %s\n", FLAGS_out.c_str());
+
+        fillAndSave(index, numTrain, num, dim);
+    } else if (FLAGS_ivfflat) {
+        faiss::IndexFlatL2 quantizerL2(dim);
+        faiss::IndexFlatIP quantizerIP(dim);
+
+        faiss::IndexFlat* quantizer = FLAGS_l2
+                ? (faiss::IndexFlat*)&quantizerL2
+                : (faiss::IndexFlat*)&quantizerIP;
+
+        faiss::IndexIVFFlat index(
+                quantizer,
+                dim,
+                numCentroids,
+                FLAGS_l2 ? faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT);
+
+        printf("IVFFlat: metric %s\n", FLAGS_l2 ? "L2" : "IP");
+        printf("Lists: %d\n", numCentroids);
+        printf("Database: dim %d num vecs %d trained on %d\n",
+               dim,
+               num,
+               numTrain);
+        printf("output file: %s\n", FLAGS_out.c_str());
+
+        fillAndSave(index, numTrain, num, dim);
+    }
+
+    return 0;
 }