faiss

Author: shawroad

Faiss的使用/001-使用欧式距离和内积的方式检索.py

@@ -0,0 +1,67 @@
+"""
+@file   : 001-欧式距离检索.py
+@author : xiaolu
+@email  : luxiaonlp@163.com
+@time   : 2021-08-09
+"""
+import faiss
+import numpy as np
+
+
+if __name__ == '__main__':
+    n_data, d = 1000, 512   # 检索库中的向量个数, 每个向量的维度
+    np.random.seed(43)   # 随机种子 为了多次执行结果一致
+
+    # 检索库的构造
+    data = []
+    mu, sigma = 3, 0.1   # 这里时通过高斯分布随机产生若干向量，这两个参数为均值和方差
+    for i in range(n_data):
+        data.append(np.random.normal(mu, sigma, d))
+    data = np.array(data).astype('float32')   # faiss只支持32位的浮点数
+
+    # 检索向量的生成
+    query = []
+    n_query = 10   # 生成10个query向量
+    mu, sigma = 3, 0.1
+    np.random.seed(12)
+    for i in range(n_query):
+        query.append(np.random.normal(mu, sigma, d))
+    query = np.array(query).astype('float32')
+
+    # 构建索引  记住要传入向量维度d
+    index = faiss.IndexFlatL2(d)
+    # print(index.is_trained)    # 这里若是false就要训练  后面讲
+
+    # 添加数据
+    index.add(data)
+    # print(index.ntotal)   # 总的数据量
+
+    # 开始检索
+    k = 10   # 指定让其返回10个距离最近的
+
+    # 这里我们选取data中的前五个 容易看到结果，因为自己跟自己距离肯定为0 所以最相关的肯定是自己
+    query_self = data[:5]
+
+    dis, ind = index.search(query_self, k=k)
+    print(dis)   # 每条数据代表了当前这个query 与最相关的十个数据的距离
+    print(ind)   # 每条数据代表了当前这个query 最相关的十条数据的索引
+    """
+    [[0.        8.55197   8.634906  8.683499  8.698736  8.821949  8.902446
+      8.943979  8.9516735 8.972908 ]
+     [0.        8.369204  8.482748  8.53028   8.581224  8.680499  8.684254
+      8.697291  8.719812  8.753435 ]
+     [0.        8.209936  8.392483  8.456179  8.473589  8.480727  8.551348
+      8.553277  8.576391  8.592704 ]
+     [0.        8.473689  8.621014  8.827385  8.883725  8.980131  8.99064
+      9.015673  9.017438  9.027972 ]
+     [0.        8.268832  8.349455  8.597895  8.611757  8.658188  8.675722
+      8.685029  8.70588   8.707612 ]]
+    [[  0 877 502  42 606 366 348 923 563  56]
+     [  1 849 974 106 348 364 877 242 280 173]
+     [  2 877 127 655 253 233 558 678  13 208]
+     [  3 421  94 348 502 402 536 646 563 735]
+     [  4 986 230 209 446 889 974 241 550 248]]
+     """
+
+
+

Faiss的使用/001-欧式距离检索.py

@@ -0,0 +1,42 @@
+"""
+@file   : 002-倒排表快速索引.py
+@author : xiaolu
+@email  : luxiaonlp@163.com
+@time   : 2021-08-09
+"""
+import numpy as np
+import faiss
+
+if __name__ == '__main__':
+    n_data, d = 1000, 512  # 检索库中的向量个数, 每个向量的维度
+    np.random.seed(43)  # 随机种子 为了多次执行结果一致
+
+    # 检索库的构造
+    data = []
+    mu, sigma = 3, 0.1  # 这里时通过高斯分布随机产生若干向量，这两个参数为均值和方差
+    for i in range(n_data):
+        data.append(np.random.normal(mu, sigma, d))
+    data = np.array(data).astype('float32')  # faiss只支持32位的浮点数
+
+    # 检索向量的生成
+    query = []
+    n_query = 10  # 生成10个query向量
+    mu, sigma = 3, 0.1
+    np.random.seed(12)
+    for i in range(n_query):
+        query.append(np.random.normal(mu, sigma, d))
+    query = np.array(query).astype('float32')
+
+    nlist = 50  # 将数据库向量分割为多少了维诺空间
+    k = 10
+    quantizer = faiss.IndexFlatL2(d)  # 量化器
+    index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)  # METRIC_L2计算L2距离, 或faiss.METRIC_INNER_PRODUCT计算内积
+    assert not index.is_trained  # 倒排表索引类型需要训练
+    index.train(data)  # 训练数据集应该与数据库数据集同分布
+    assert index.is_trained
+
+    index.add(data)
+    index.nprobe = 2  # 选择n个维诺空间进行索引,
+    dis, ind = index.search(query, k)
+    print(dis)
+    print(ind)

Faiss的使用/002-IndexIVFFlat检索.py

@@ -0,0 +1,45 @@
+"""
+@file   : 003-乘积量化索引.py
+@author : xiaolu
+@email  : luxiaonlp@163.com
+@time   : 2021-08-09
+"""
+import numpy as np
+import faiss
+
+if __name__ == '__main__':
+    n_data, d = 1000, 512  # 检索库中的向量个数, 每个向量的维度
+    np.random.seed(43)  # 随机种子 为了多次执行结果一致
+
+    # 检索库的构造
+    data = []
+    mu, sigma = 3, 0.1  # 这里时通过高斯分布随机产生若干向量，这两个参数为均值和方差
+    for i in range(n_data):
+        data.append(np.random.normal(mu, sigma, d))
+    data = np.array(data).astype('float32')  # faiss只支持32位的浮点数
+
+    # 检索向量的生成
+    query = []
+    n_query = 10  # 生成10个query向量
+    mu, sigma = 3, 0.1
+    np.random.seed(12)
+    for i in range(n_query):
+        query.append(np.random.normal(mu, sigma, d))
+    query = np.array(query).astype('float32')
+
+    nlist = 50
+    m = 8  # 列方向划分个数，必须能被d整除
+    k = 10
+    quantizer = faiss.IndexFlatL2(d)
+    index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 4)   # 4 表示每个子向量被编码为 4 bits
+
+    index.train(data)
+    index.add(data)
+    index.nprobe = 50
+    dis, ind = index.search(data[:10], k)  # 查询自身
+    print(dis)
+    print(ind)
+
+    dis, ind = index.search(query, k)  # 真实查询
+    print(dis)
+    print(ind)

Faiss的使用/002-倒排表快速索引.py

@@ -0,0 +1,58 @@
+"""
+@file   : 004-faiss实现kmeans聚类.py
+@author : xiaolu
+@email  : luxiaonlp@163.com
+@time   : 2021-08-09
+"""
+import faiss
+import numpy as np
+
+
+if __name__ == '__main__':
+    # 数据
+    n_data, d = 2000, 512
+    np.random.seed(43)
+    data = []
+    mu, sigma = 3, 0.1
+    for i in range(n_data):
+        data.append(np.random.normal(mu, sigma, d))
+    data = np.array(data).astype('float32')
+
+    # 聚类
+    n_centroids = 1024    # 聚类中心个数
+    d = data.shape[1]
+    kmeans = faiss.Kmeans(d, n_centroids)
+    kmeans.train(data)
+    # 输出聚类中心
+    # print(kmeans.centroids)
+    # print(len(kmeans.centroids))
+
+    # 看data中的前五个向量属于那个类(最有可能的两个类)
+    D, I = kmeans.index.search(data[:5], k=2)
+    print(D)   # 与每个类的距离
+    print(I)   # 类的编号
+    """
+    输出:
+    [[4.1553707 5.2924204]
+     [1.9329664 4.930997 ]
+     [4.537619  4.8509283]
+     [4.6700296 5.2252126]
+     [2.101182  4.9292693]]
+    [[478 568]
+     [767 697]
+     [568 527]
+     [999 568]
+     [175 853]]
+    """
+
+    print('*'*100)
+    # 计算每个中心最近的若干条向量
+    k = 5
+    index = faiss.IndexFlatL2(d)
+    index.add(data)
+    D, I = index.search(kmeans.centroids, k)
+    print(D)
+    print(I)
+
+
+

Faiss的使用/003-乘积量化索引.py

@@ -0,0 +1,27 @@
+"""
+@file   : 005-faiss实现pca降维.py
+@author : xiaolu
+@email  : luxiaonlp@163.com
+@time   : 2021-08-09
+"""
+import faiss
+import numpy as np
+
+
+if __name__ == '__main__':
+    # 数据
+    n_data, d = 2000, 512
+    np.random.seed(43)
+    data = []
+    mu, sigma = 3, 0.1
+    for i in range(n_data):
+        data.append(np.random.normal(mu, sigma, d))
+    data = np.array(data).astype('float32')
+
+    mat = faiss.PCAMatrix(512, 64)  # 从512维降为64维
+    mat.train(data)
+    assert mat.is_trained
+    tr = mat.apply_py(data)
+    print(tr.shape)
+    print(tr)
+

Faiss的使用/004-faiss实现kmeans聚类.py

@@ -0,0 +1,37 @@
+"""
+@file   : 006-faiss实现PQ编码和解码.py
+@author : xiaolu
+@email  : luxiaonlp@163.com
+@time   : 2021-08-09
+"""
+import faiss
+import numpy as np
+
+
+if __name__ == '__main__':
+    # 数据
+    n_data, d = 2000, 512
+    np.random.seed(43)
+    data = []
+    mu, sigma = 3, 0.1
+    for i in range(n_data):
+        data.append(np.random.normal(mu, sigma, d))
+    data = np.array(data).astype('float32')
+
+    cs = 4  # code size (bytes)
+    # 训练数据集
+    x = data   # 原始的数据集
+
+    x_train = data  # 训练集
+    pq = faiss.ProductQuantizer(d, cs, 8)
+    pq.train(x_train)
+
+    # encode编码
+    codes = pq.compute_codes(x)
+
+    # decode解码
+    x2 = pq.decode(codes)
+
+    # 编码-解码后与原始数据的差
+    avg_relative_error = ((x - x2)**2).sum() / (x ** 2).sum()
+    print(avg_relative_error)