乘积量化(PQ)算法的nanopq库python

最新推荐文章于 2024-10-18 15:44:13 发布

nangua06

最新推荐文章于 2024-10-18 15:44:13 发布

阅读量739

点赞数 1

文章标签： python

本文链接：https://blog.youkuaiyun.com/nangua06/article/details/112965139

版权

这段代码展示了如何使用nanopq库来训练PQ编码，并对高维数据进行编码和检索。首先，生成了随机的训练数据和查询向量，然后通过PQ编码将训练数据转换为PQ码。接着，利用PQ编码和距离表进行在线检索，计算查询向量与数据集之间的Asymmetric Distance。最后，将码字保存到文件并输出了PQ码和距离表的大小信息。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import nanopq
import numpy as np

N, Nt, D = 1000000, 100000, 128
X = np.random.random((N, D)).astype(np.float32)  # 10,000 128-dim vectors to be indexed

Xt = np.random.random((Nt, D)).astype(np.float32)  # 2,000 128-dim vectors for training
print('Xt')
print(Xt.shape)

query = np.random.random((D,)).astype(np.float32)  # a 128-dim query vector

#print(query)
# Instantiate with M=8 sub-spaces
pq = nanopq.PQ(M=16)

# Train codewords
pq.fit(Xt)

print('pq')
print(pq.codewords)

with open('/Users/dl/Documents/bak/' + "codewords.txt", 'w') as fp:
    fp.write(str(pq.codewords))
print('codewords')
print(np.array(pq.codewords).shape)

# Encode to PQ-codes
X_code = pq.encode(X)  # (10000, 8) with dtype=np.uint8
print('X_code')
print(X_code.shape)
print(X_code)

# Results: create a distance table online, and compute Asymmetric Distance to each PQ-code
dists = pq.dtable(query).adist(X_code)

dtable = pq.dtable(query)
print('dtable')
print(dtable.__sizeof__())
print(dtable)
print('dists')
print(dists.shape)
print(dists)