Updata export_index

This commit is contained in:
ylzz1997 2023-07-23 23:55:42 +08:00
parent b306a0ecfa
commit 121205c090
1 changed files with 20 additions and 0 deletions

20
export_index_for_onnx.py Normal file
View File

@ -0,0 +1,20 @@
import os
import pickle
import faiss
path = "crs"
indexs_file_path = f"checkpoints/{path}/feature_and_index.pkl"
indexs_out_dir = f"checkpoints/{path}/"
with open("feature_and_index.pkl",mode="rb") as f:
indexs = pickle.load(f)
for k in indexs:
print(f"Save {k} index")
faiss.write_index(
indexs[k],
os.path.join(indexs_out_dir,f"Index-{k}.index")
)
print("Saved all index")