Spaces:
Runtime error
Runtime error
File size: 2,919 Bytes
a8c39f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import os
import sys
import faiss
import numpy as np
from sklearn.cluster import MiniBatchKMeans
from multiprocessing import cpu_count
# Parse command line arguments
exp_dir = str(sys.argv[1])
version = str(sys.argv[2])
index_algorithm = str(sys.argv[3])
try:
feature_dir = os.path.join(exp_dir, f"{version}_extracted")
model_name = os.path.basename(exp_dir)
index_filename_added = f"added_{model_name}_{version}.index"
index_filepath_added = os.path.join(exp_dir, index_filename_added)
# index_filename_trained = f"trained_{model_name}_{version}.index"
# index_filepath_trained = os.path.join(exp_dir, index_filename_trained)
if os.path.exists(index_filepath_added):
pass
else:
npys = []
listdir_res = sorted(os.listdir(feature_dir))
for name in listdir_res:
file_path = os.path.join(feature_dir, name)
phone = np.load(file_path)
npys.append(phone)
big_npy = np.concatenate(npys, axis=0)
big_npy_idx = np.arange(big_npy.shape[0])
np.random.shuffle(big_npy_idx)
big_npy = big_npy[big_npy_idx]
if big_npy.shape[0] > 2e5 and (
index_algorithm == "Auto" or index_algorithm == "KMeans"
):
big_npy = (
MiniBatchKMeans(
n_clusters=10000,
verbose=True,
batch_size=256 * cpu_count(),
compute_labels=False,
init="random",
)
.fit(big_npy)
.cluster_centers_
)
# np.save(os.path.join(exp_dir, "total_fea.npy"), big_npy)
n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
"""
# index_trained
index_trained = faiss.index_factory(
256 if version == "v1" else 768, f"IVF{n_ivf},Flat"
)
index_ivf_trained = faiss.extract_index_ivf(index_trained)
index_ivf_trained.nprobe = 1
index_trained.train(big_npy)
faiss.write_index(index_trained, index_filepath_trained)
"""
# index_added
index_added = faiss.index_factory(
256 if version == "v1" else 768, f"IVF{n_ivf},Flat"
)
index_ivf_added = faiss.extract_index_ivf(index_added)
index_ivf_added.nprobe = 1
index_added.train(big_npy)
batch_size_add = 8192
for i in range(0, big_npy.shape[0], batch_size_add):
index_added.add(big_npy[i : i + batch_size_add])
faiss.write_index(index_added, index_filepath_added)
print(f"Saved index file '{index_filepath_added}'")
except Exception as error:
print(f"An error occurred extracting the index: {error}")
print(
"If you are running this code in a virtual environment, make sure you have enough GPU available to generate the Index file."
)
|