mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Michael Norris
4030faae3d
Summary: This PR adds AVX-512 implementations of the four hot functions in polysemous training (`compute_cost` and `cost_update` for both `ReproduceWithHammingObjective` and `ReproduceDistancesObjective`), integrated via FAISS's SIMD dynamic dispatch framework. It speeds up the training phase by up to **1.09x**. ### Benchmarks Training phase of `benchs/bench_polysemous_sift1m.py` on Sapphire Rapids (SPR): ```bash $ numactl -m 0 -C 0-7 python benchs/bench_polysemous_sift1m.py ``` | Build | Median training time | Speedup | |-------|---------------------|---------| | Scalar (baseline) | ~4.29 s | 1.00x | | AVX-512 (this PR) | ~3.95 s | 1.09x | Search accuracy and latency are unchanged — the optimization only affects the training path. Reviewed By: mdouze Differential Revision: D106200518 fbshipit-source-id: d0037ef29c3eababf0b6d15f37c03a05b13a7cd5
68 lines
1.6 KiB
Python
68 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from __future__ import print_function
|
|
|
|
import time
|
|
import numpy as np
|
|
import faiss
|
|
from datasets import load_sift1M, evaluate
|
|
|
|
NUM_TRAIN_RUNS = 5
|
|
|
|
print("load data")
|
|
xb, xq, xt, gt = load_sift1M()
|
|
nq, d = xq.shape
|
|
|
|
train_times = []
|
|
for run in range(NUM_TRAIN_RUNS):
|
|
index = faiss.IndexPQ(d, 16, 8)
|
|
index.do_polysemous_training = True
|
|
index.verbose = run == 0
|
|
|
|
print("train run %d/%d" % (run + 1, NUM_TRAIN_RUNS))
|
|
|
|
t0 = time.time()
|
|
index.train(xt)
|
|
t1 = time.time()
|
|
elapsed = t1 - t0
|
|
train_times.append(elapsed)
|
|
print(" Training time: %.2f s" % elapsed)
|
|
|
|
times = np.array(train_times)
|
|
print(
|
|
"\nTraining time over %d runs: "
|
|
"median %.2f s, mean %.2f s, std %.2f s, min %.2f s, max %.2f s"
|
|
% (
|
|
NUM_TRAIN_RUNS,
|
|
np.median(times),
|
|
np.mean(times),
|
|
np.std(times),
|
|
np.min(times),
|
|
np.max(times),
|
|
)
|
|
)
|
|
|
|
print("\nadd vectors to index")
|
|
|
|
index.add(xb)
|
|
|
|
nt = 1
|
|
faiss.omp_set_num_threads(1)
|
|
|
|
|
|
print("PQ baseline", end=" ")
|
|
index.search_type = faiss.IndexPQ.ST_PQ
|
|
t, r = evaluate(index, xq, gt, 1)
|
|
print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1]))
|
|
|
|
for ht in 64, 62, 58, 54, 50, 46, 42, 38, 34, 30:
|
|
print("Polysemous", ht, end=" ")
|
|
index.search_type = faiss.IndexPQ.ST_polysemous
|
|
index.polysemous_ht = ht
|
|
t, r = evaluate(index, xq, gt, 1)
|
|
print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1]))
|