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