Allow k and M suffixes in IVF indexes (#3812)

Summary:
Pull Request resolved: #3812

Allows factory strings like `IVF3k,Flat` as a shorthand for 3072 centroids.

The main question is whether k or M should be metric (k=1000) or power of 2 (k=1024):

* pro-metric: standard,

* pro-power of 2: in practice we use powers of 2 most often

The suffixes ki and Mi should be used for powers of 2 but this makes the notation more heavy (which is what we wanted to avoid in the first place).

So I picked power of 2.

Reviewed By: mnorris11

Differential Revision: D62019941

fbshipit-source-id: f547962625123ecdfaa406067781c77386017793

faiss/index_factory.cpp

@@ -226,6 +226,19 @@ VectorTransform* parse_VectorTransform(const std::string& description, int d) {
  * Parse IndexIVF
  */
 
+size_t parse_nlist(std::string s) {
+ size_t multiplier = 1;
+ if (s.back() == 'k') {
+ s.pop_back();
+ multiplier = 1024;
+ }
+ if (s.back() == 'M') {
+ s.pop_back();
+ multiplier = 1024 * 1024;
+ }
+ return std::stoi(s) * multiplier;
+}
+
 // parsing guard + function
 Index* parse_coarse_quantizer(
  const std::string& description,
@@ -240,8 +253,8 @@ Index* parse_coarse_quantizer(
  };
  use_2layer = false;
 
- if (match("IVF([0-9]+)")) {
- nlist = std::stoi(sm[1].str());
+ if (match("IVF([0-9]+[kM]?)")) {
+ nlist = parse_nlist(sm[1].str());
  return new IndexFlat(d, mt);
  }
  if (match("IMI2x([0-9]+)")) {
@@ -252,18 +265,18 @@ Index* parse_coarse_quantizer(
  nlist = (size_t)1 << (2 * nbit);
  return new MultiIndexQuantizer(d, 2, nbit);
  }
- if (match("IVF([0-9]+)_HNSW([0-9]*)")) {
- nlist = std::stoi(sm[1].str());
+ if (match("IVF([0-9]+[kM]?)_HNSW([0-9]*)")) {
+ nlist = parse_nlist(sm[1].str());
  int hnsw_M = sm[2].length() > 0 ? std::stoi(sm[2]) : 32;
  return new IndexHNSWFlat(d, hnsw_M, mt);
  }
- if (match("IVF([0-9]+)_NSG([0-9]+)")) {
- nlist = std::stoi(sm[1].str());
+ if (match("IVF([0-9]+[kM]?)_NSG([0-9]+)")) {
+ nlist = parse_nlist(sm[1].str());
  int R = std::stoi(sm[2]);
  return new IndexNSGFlat(d, R, mt);
  }
- if (match("IVF([0-9]+)\\(Index([0-9])\\)")) {
- nlist = std::stoi(sm[1].str());
+ if (match("IVF([0-9]+[kM]?)\\(Index([0-9])\\)")) {
+ nlist = parse_nlist(sm[1].str());
  int no = std::stoi(sm[2].str());
  FAISS_ASSERT(no >= 0 && no < parenthesis_indexes.size());
  return parenthesis_indexes[no].release();

tests/test_factory.py

@@ -238,6 +238,18 @@ def test_ivf(self):
  index = faiss.index_factory(123, "IVF456,Flat")
  self.assertEqual(index.__class__, faiss.IndexIVFFlat)
 
+ def test_ivf_suffix_k(self):
+ index = faiss.index_factory(123, "IVF3k,Flat")
+ self.assertEqual(index.nlist, 3072)
+
+ def test_ivf_suffix_M(self):
+ index = faiss.index_factory(123, "IVF1M,Flat")
+ self.assertEqual(index.nlist, 1024 * 1024)
+
+ def test_ivf_suffix_HNSW_M(self):
+ index = faiss.index_factory(123, "IVF1M_HNSW,Flat")
+ self.assertEqual(index.nlist, 1024 * 1024)
+
  def test_idmap(self):
  index = faiss.index_factory(123, "Flat,IDMap")
  self.assertEqual(index.__class__, faiss.IndexIDMap)