nmslib完成大规模高维向量查找（如大规模人脸识别）

最近在完成一个人脸识别项目，标准人脸库里面大约有50万张人脸图像，人脸识别算法使用的是MTCNN + ArcFace，每张人脸图片经过ArcFace模型转换后形成一个512维的特征向量，那么50万张人脸图像形成的人脸特征矩阵大小就是[500000, 512]。

在进行人脸识别时，通过计算待识别的人脸特征与这50万个特征的最小距离来确定人物身份。对于每次只识别一个人脸特征来说，使用numpy的广播机制，可以完成最小距离的计算。但是现在如果要识别图片上的10个人脸，这10张人脸的特征矩阵是[10, 512]，这个矩阵无法通过numpy的广播机制与[500000， 512]矩阵进行广播运算。那么只能使用for循环对10个人脸逐个进行识别，这种情况下识别一张图片消耗的时间就与图片上的人脸个数是完全线性的关系。这显然是无法接受的。

---

下面简单介绍下如何使用nmslib完成大规模高维向量的计算，nmslib使用过程中具体的参数配置详见nmslib开源代码，本篇内容提供的配置参数对于几百万的高维向量查找足够了。

1、nmslib安装

在python环境下安装nmslib很很简单，pip install nmslib就可以完成，完后程使用import nmslib进行导入即可

2、nmslib具体使用示例

# coding:utf-8

import numpy as np
import nmslib
import datetime
from functools import wraps


def func_execute_time(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = datetime.datetime.now()
        res = func(*args, **kwargs)
        end_time = datetime.datetime.now()
        duration_time = (end_time - start_time).microseconds // 1000
        print("execute function %s, elapse time %.4f ms" % (func.__name__, duration_time))
        return res
    return wrapper


def create_indexer(vec, index_thread, m, ef):
    """
    基于数据向量构建索引
    :param vec: 原始数据向量
    :param index_thread: 线程数
    :param m: 参照官网
    :param ef: 参照官网
    :return:
    """
    index = nmslib.init(method="hnsw", space="l2")
    index.addDataPointBatch(vec)
    INDEX_TIME_PARAMS = {
        "indexThreadQty": index_thread,
        "M": m,
        "efConstruction": ef,
        "post": 2
    }
    index.createIndex(INDEX_TIME_PARAMS, print_progress=True)
    index.saveIndex("data_%d_%d_%d.hnsw" % (index_thread, m, ef))


def load_indexer(index_path, ef_search=300):
    """
    加载构建好的向量索引文件
    :param index_path: 索引文件地址
    :param ef_search: 查询结果参数
    :return:
    """
    indexer = nmslib.init(method="hnsw", space="l2")
    indexer.loadIndex(index_path)
    indexer.setQueryTimeParams({"efSearch": ef_search})
    return indexer


@func_execute_time
def search_vec_top_n(indexer, vecs, top_n=7, threads=4):
    """
    使用构建好的索引文件完成向量查询
    :param indexer: 索引
    :param vecs: 待查询向量
    :param top_n: 返回前top_n个查询结果
    :param threads:
    :return:
    """
    neighbours = indexer.knnQueryBatch(vecs, k=top_n, num_threads=threads)
    # print(neighbours)
    return neighbours


if __name__ == '__main__':
    data = np.arange(300000000).reshape(1000000, 300)
    print(data.shape)
    create_indexer(data, 10, 5, 300)
    indexer = load_indexer("data_10_5_300.hnsw")

    print(search_vec_top_n(indexer, data[:10]))
    print(search_vec_top_n(indexer, data[-10:]))

1、上述代码中构建了一个大小为[1000000, 300]的矩阵，并使用create_indexer函数构建矩阵索引，索引文件保存在*.hnsw文件中

构建过程中涉及到的参数：

method="hnsw"，全称叫“Hierarchical Navigable Small World Graph”是nmslib中使用的一种索引构建算法。

space="l2"，指定查找向量时使用的距离度量方法

indexThreadQty、M、efConstruction、post等参数用来权衡构建的索引复杂度和最终的查找精度上的权衡，对于几百万个向量的查找，参照上述参数设置即可。具体参数含义参考nmslib手册

2、对于构建好的索引文件使用load_indexer进行加载

加载索引文件过程中涉及到的参数：efSearch的设置保持与efConstruction一致

3、在search_vec_top_n函数中使用构建好的所以完成向量查找

查找阶段涉及的参数：

top_n在向量空间中查找最近的top_n个向量

thread指定查找时的线程数，与查找速度相关

---

上述代码运行结果如下：

data[:10]向量的查找结果为：

execute function search_vec_top_n, elapse time 15.0000 ms
[
    (array([0, 1, 2, 3, 4, 5, 6]), array([    0.    ,  5196.1523, 10392.305 , 15588.457 , 20784.61  , 25980.762 , 31176.914 ], dtype=float32)), 
    (array([1, 0, 2, 3, 4, 5, 6]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 15588.457 , 20784.61  , 25980.762 ], dtype=float32)), 
    (array([2, 1, 3, 0, 4, 5, 6]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 20784.61  ], dtype=float32)), 
    (array([3, 2, 4, 1, 5, 0, 6]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32)), 
    (array([4, 3, 5, 2, 6, 1, 7]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32)), 
    (array([5, 4, 6, 3, 7, 2, 8]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32)), 
    (array([6, 5, 7, 4, 8, 3, 9]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32)), 
    (array([ 7,  6,  8,  5,  9,  4, 10]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32)), 
    (array([ 8,  7,  9,  6, 10,  5, 11]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32)), 
    (array([ 9,  8, 10,  7, 11,  6, 12]), array([    0.    ,  5196.1523,  5196.1523, 10392.305 , 10392.305 , 15588.457 , 15588.457 ], dtype=float32))
]

data[:-10]向量的查找结果为：

execute function search_vec_top_n, elapse time 15.0000 ms
[
    (array([999990, 999989, 999991, 999988, 999992, 999987, 999993]), array([    0.    ,  5200.271 ,  5209.6157, 10385.96  , 10395.076 , 15594.214 , 15594.214 ], dtype=float32)), 
    (array([999991, 999992, 999990, 999993, 999989, 999988, 999994]), array([    0.    ,  5196.528 ,  5209.6157, 10393.253 , 10398.72  , 15586.727 , 15586.727 ], dtype=float32)), 
    (array([999992, 999991, 999993, 999990, 999994, 999989, 999995]), array([    0.   ,  5196.528,  5207.748, 10395.076, 10398.72 , 15586.727, 15586.727], dtype=float32)), 
    (array([999993, 999994, 999992, 999995, 999991, 999990, 999996]), array([    0.   ,  5202.141,  5207.748, 10387.783, 10393.253, 15594.214, 15594.214], dtype=float32)), 
    (array([999994, 999995, 999993, 999992, 999996, 999991, 999997]), array([    0.   ,  5196.528,  5202.141, 10398.72 , 10400.542, 15586.727, 15586.727], dtype=float32)), 
    (array([999995, 999994, 999996, 999993, 999997, 999992, 999998]), array([    0.   ,  5196.528,  5215.215, 10387.783, 10398.72 , 15586.727, 15592.343], dtype=float32)), 
    (array([999996, 999997, 999995, 999998, 999994, 999999, 999993]), array([    0.   ,  5194.656,  5215.215, 10385.96 , 10400.542, 15588.599, 15594.214], dtype=float32)), 
    (array([999997, 999996, 999998, 999995, 999999, 999994, 999993]), array([    0.   ,  5194.656,  5202.141, 10398.72 , 10402.363, 15586.727, 20782.762], dtype=float32)), 
    (array([999998, 999997, 999999, 999996, 999995, 999994, 999993]), array([    0.   ,  5202.141,  5211.483, 10385.96 , 15592.343, 20782.762, 25976.826], dtype=float32)), 
    (array([999999, 999998, 999997, 999996, 999995, 999994, 999993]), array([    0.   ,  5211.483, 10402.363, 15588.599, 20797.54 , 25985.99 , 31181.549], dtype=float32))
]

data[:10]和data[:-10]分别是data向量的前10个和后10个，在l2空间距离上与他们最近的7个向量分别是他们各自为中心，以及他们左右各3个向量。比如以data[:-10]中的第一个向量（也就是data[-10]，索引编号为999990）的计算结果为例：

(array([999990, 999989, 999991, 999988, 999992, 999987, 999993]), array([    0.    ,  5200.271 ,  5209.6157, 10385.96  , 10395.076 , 15594.214 , 15594.214 ], dtype=float32))

与data[-10]距离最近的向量是他自己，距离为0.0，其次就是data[-10]的前一个向量，索引编号为999989的向量，在下一个就是data[-10]的后一个向量，索引编号为999991的向量。

在上述结果中可以看到，在[1000000, 300]大小的矩阵上查找10个大小为300维的向量话费大约15毫秒左右的时间，表明nmslib构建索引后不仅可以保证查找结果的准确性，同时大大提升了查找效率。

---

最后讲一下在人脸识别项目中nmslib的使用，首先使用nmslib构建使用算法处理好的人脸特征向量的索引文件，在构建索引文件的同时，使用一个映射文件记录每个人脸特征向量的索引编号及其对应的人物身份。如（1，小明），（2，老王），……，然后使用索引编号与nmslib查询结果中的索引编号进行关联，就可以确定人脸的身份了。