53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
|
# import numpy as np
|
||
|
|
# a = [1,2,3,4,5,6,7,8,9,10]
|
||
|
|
# print(a[0::2])
|
||
|
|
# print(a[1::2])
|
||
|
|
#
|
||
|
|
# src = np.array([
|
||
|
|
# [30.2946, 51.6963],
|
||
|
|
# [65.5318, 51.5014],
|
||
|
|
# [48.0252, 71.7366],
|
||
|
|
# [33.5493, 92.3655],
|
||
|
|
# [62.7299, 92.2041]], dtype=np.float32)
|
||
|
|
# src[:, 0] += 8.0
|
||
|
|
# #print(src[:, 1])
|
||
|
|
# input_blob = np.zeros((2, 3, 2, 2), dtype=np.uint8)
|
||
|
|
# #print(input_blob[0])
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
d = 512 # dimension
|
||
|
|
nb = 23800 # database size
|
||
|
|
nq = 1 # nb of queries
|
||
|
|
np.random.seed(1234) # make reproducible
|
||
|
|
xb = np.random.random((nb, d)).astype('float32')
|
||
|
|
#print(xb)
|
||
|
|
#print(np.arange(nb))
|
||
|
|
#print(xb[:, 0])
|
||
|
|
|
||
|
|
# xb[:, 0] += np.arange(nb) / 1000.
|
||
|
|
# #print(xb)
|
||
|
|
#xq = np.random.random((nq, d)).astype('float32')
|
||
|
|
#xq[:, 0] += np.arange(nq) / 1000.
|
||
|
|
import faiss # make faiss available
|
||
|
|
nlist = 100
|
||
|
|
#quantizer = faiss.IndexFlatL2(d) # the other index
|
||
|
|
#index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
|
||
|
|
#index.train(xb)
|
||
|
|
index = faiss.IndexFlatL2(d) # build the index
|
||
|
|
print(index.is_trained)
|
||
|
|
#print(xq.shape)
|
||
|
|
index.add(xb) # add vectors to the index
|
||
|
|
print(index.ntotal)
|
||
|
|
#index.nprobe=30
|
||
|
|
k = 1 # we want to see 4 nearest neighbors
|
||
|
|
import time
|
||
|
|
for i in range(10):
|
||
|
|
xq = np.random.random((nq, d)).astype('float32')
|
||
|
|
start_time = time.time()
|
||
|
|
D, I = index.search(xq, k) # actual search
|
||
|
|
end_time = time.time()
|
||
|
|
print("faiss cost %fs"%(end_time - start_time))
|
||
|
|
print(D,I)
|
||
|
|
|
||
|
|
|