import numpy as np
from scipy.cluster.vq import vq, kmeans2
from scipy.spatial.distance import cdist
def train(vec, M, Ks):
"""
:param vec: 向量
:param M: 子向量组数
:param Ks: 每组向量聚类个数
:return: codeword: [M, Ks, Ds],
codeword[m][k]表示第m组子向量第k个子向量所属的聚类中心向量
"""
Ds = int(vec.shape[1] / M)
codeword = np.empty((M, Ks, Ds), np.float32)
for m in range(M):
vec_sub = vec[:, m * Ds: (m + 1) * Ds]
# 第m组子向量vec_sub聚成Ks类
# kmeans2返回两个结果,第一个是原始向量归属类目的中心向量,第二个是类目ID
codeword[m], label = kmeans2(vec_sub, Ks)
return codeword
def encode(codeword, vec):
"""
:param codeword: 码本,shape为[M, Ks, Ds]
:param vec: 原始向量
:return: pqcode: pq编码结果,
shape为[N, M],每个原始向量用M组子向量的聚类中心ID表示
"""
M, Ks, Ds = codeword.shape
# pq编码shape为[N, M]
pqcode = np.empty((vec.shape[0], M), np.int64)
for m in range(M):
vec_sub = vec[:, m * Ds: (m + 1) * Ds]
# 第m组子向量
# 第m组子向量中每个子向量在第m个码本中查找距离最近的
pqcode[:, m], dist = vq(vec_sub, codeword[m])
return pqcode
def search(codeword, pqcode, query):
"""
:param codeword:
:param pqcode: pq编码结果, shape为[N, M],每个原始向量用M组子向量的聚类中心ID表示
:param query: 查询向量[1, d]
:return: dist:查询向量与原始向量的距离,shape为[N,]
"""
M, Ks, Ds = codeword.shape
# 距离向量表, [M, Ks]
dist_table = np.empty((M, Ks))
for m in range(M):
query_sub = query[m * Ds: (m + 1) * Ds]
# query_sub向量与第m个码本每个向量距离
dist_table[m, :] = cdist([query_sub], codeword[m], 'sqeuclidean')[0]
# dist_table[range(M), pqcode] 为 query向量与原始向量在每个子向量的聚类,shape为[N, M]
# 每组子向量距离相加
dist = np.sum(dist_table[range(M), pqcode], axis=1)
return dist
def main():
# 数据量
N = 50000
# 向量维度
d = 128
# 每组子向量聚类个数
Ks = 32
# 训练向量[N, d]
vec_train = np.random.random((N, d))
# 查询向量[1, d]
# mock 第100个是距离查询向量最近的
selected_vec = vec_train[100]
query_vec = selected_vec + [np.random.uniform(-0.001, 0.001) for _ in range(d)]
query = np.random.random((1, d))
# 子向量组数
M = 4
# 对原始向量划分子向量组,并对每组子向量进行聚类
codeword = train(vec_train, M, Ks)
# pq编码
pqcode = encode(codeword, vec_train)
# 查询向量
dist = search(codeword, pqcode, query_vec)
sorted_dist = sorted(enumerate(dist), key=lambda x: x[1])
print(sorted_dist[0])
"""
(100, 6.850794458722508)
"""
# -*- coding:utf-8 -*-
import faiss
import numpy as np
def test_IndexFlatL2(vec_train, query, top_k=5):
"""
暴力检索
:param vec_train:
:param query:
:return:
"""
N, d = vec_train.shape
# 1. 创建索引
index = faiss.IndexFlatL2(d)
# 2. 添加数据集
index.add(vec_train)
# 3. 检索
dist_list, label_list = index.search(np.array([query]), k=top_k)
print(dist_list, label_list)
def test_IndexIVFFlat(vec_train, query_vec, top_k=5):
"""
通过创建倒排索引优化
流程:
使用k-means对train向量进行聚类,查询时query_vec所归属的类目中进行检索
:param vec_train:
:param query_vec:
:param top_k:
:return:
"""
nlist = 100 # 聚类中心的个数
N, d = vec_train.shape
quantizer = faiss.IndexFlatL2(d) # the other index
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
# 添加 训练集
index.train(vec_train)
index.add(vec_train)
# 检索
res = index.search(query_vec, k=top_k)
print(res)
def main():
# 数据量
N = 50000
# 向量维度
d = 128
vec_train = np.ascontiguousarray(np.random.random((N, d)), np.float32)
# mock 第100个是距离查询向量最近的
selected_vec = vec_train[100]
query_vec = selected_vec + [np.random.uniform(-0.001, 0.001) for _ in range(d)]
query_vec = np.ascontiguousarray(query_vec, np.float32)
# 1. 暴力检索,全量检索
# test_IndexFlatL2(vec_train, query_vec)
# 2. 倒排索引
test_IndexIVFFlat(vec_train, query_vec)
if __name__ == '__main__':
main()
参考
[1] 实例理解product quantization算法
[2] 【关于 Faiss 】 那些的你不知道的事-技术圈
[3] Home · facebookresearch/faiss Wiki
[4] A Survey of Product Quantization
[5] Product quantization for nearest neighbor search
[6] 理解 product quantization 算法
[7] https://github.com/matsui528/na
转子:https://zhuanlan.zhihu.com/p/534004381