1#!/usr/bin/env python3
2
3# Copyright (c) Facebook, Inc. and its affiliates.
4#
5# This source code is licensed under the MIT license found in the
6# LICENSE file in the root directory of this source tree.
7
8from __future__ import print_function
9import time
10import numpy as np
11
12import faiss
13from datasets import load_sift1M, evaluate
14
15
16print("load data")
17xb, xq, xt, gt = load_sift1M()
18nq, d = xq.shape
19
20# index with 16 subquantizers, 8 bit each
21index = faiss.IndexPQ(d, 16, 8)
22index.do_polysemous_training = True
23index.verbose = True
24
25print("train")
26
27index.train(xt)
28
29print("add vectors to index")
30
31index.add(xb)
32
33nt = 1
34faiss.omp_set_num_threads(1)
35
36
37print("PQ baseline", end=' ')
38index.search_type = faiss.IndexPQ.ST_PQ
39t, r = evaluate(index, xq, gt, 1)
40print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1]))
41
42for ht in 64, 62, 58, 54, 50, 46, 42, 38, 34, 30:
43    print("Polysemous", ht, end=' ')
44    index.search_type = faiss.IndexPQ.ST_polysemous
45    index.polysemous_ht = ht
46    t, r = evaluate(index, xq, gt, 1)
47    print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1]))
48